diff --git a/devito/data/allocators.py b/devito/data/allocators.py index 72289c57bf..aff28ef108 100644 --- a/devito/data/allocators.py +++ b/devito/data/allocators.py @@ -337,10 +337,10 @@ def put_local(self): return self._node == 'local' -class ExternalAllocator(MemoryAllocator): +class DataReference(MemoryAllocator): """ - An ExternalAllocator is used to assign pre-existing user data to Functions. + A DataReference is used to assign pre-existing user data to Functions. Thus, Devito does not allocate any memory. Parameters @@ -350,23 +350,23 @@ class ExternalAllocator(MemoryAllocator): Notes ------- - * Use ExternalAllocator and pass a reference to the external memory when + * Use DataReference and pass a reference to the external memory when creating a Function. This Function will now use this memory as its f.data. - * If the data present in this external memory is valuable, provide a noop - initialiser, or else Devito will reset it to 0. + * This can be used to pass one Function's data to another to avoid copying + during Function rebuilds (this should only be used internally). Example -------- >>> from devito import Grid, Function - >>> from devito.data.allocators import ExternalAllocator + >>> from devito.data.allocators import DataReference >>> import numpy as np >>> shape = (2, 2) >>> numpy_array = np.ones(shape, dtype=np.float32) >>> g = Grid(shape) >>> space_order = 0 >>> f = Function(name='f', grid=g, space_order=space_order, - ... allocator=ExternalAllocator(numpy_array), initializer=lambda x: None) + ... allocator=DataReference(numpy_array)) >>> f.data[0, 1] = 2 >>> numpy_array array([[1., 2.], @@ -387,6 +387,9 @@ def alloc(self, shape, dtype, padding=0): return (self.numpy_array, None) +# For backward compatibility +ExternalAllocator = DataReference + ALLOC_GUARD = GuardAllocator(1048576) ALLOC_ALIGNED = PosixAllocator() ALLOC_KNL_DRAM = NumaAllocator(0) diff --git a/devito/types/basic.py b/devito/types/basic.py index 10e07087b2..75befaf2cd 100644 --- a/devito/types/basic.py +++ b/devito/types/basic.py @@ -880,6 +880,14 @@ def __new__(cls, *args, **kwargs): # let's just return `function` itself return function + # If dimensions have been replaced, then it is necessary to set `function` + # to None. It is also necessary to remove halo and padding kwargs so that + # they are rebuilt with the new dimensions + if function is not None and function.dimensions != dimensions: + function = kwargs['function'] = None + kwargs.pop('padding', None) + kwargs.pop('halo', None) + with sympy_mutex: # Go straight through Basic, thus bypassing caching and machinery # in sympy.Application/Function that isn't really necessary diff --git a/devito/types/dense.py b/devito/types/dense.py index 4ae37269cb..d96cc9b838 100644 --- a/devito/types/dense.py +++ b/devito/types/dense.py @@ -11,6 +11,7 @@ from devito.builtins import assign from devito.data import (DOMAIN, OWNED, HALO, NOPAD, FULL, LEFT, CENTER, RIGHT, Data, default_allocator) +from devito.data.allocators import DataReference from devito.exceptions import InvalidArgument from devito.logger import debug, warning from devito.mpi import MPI @@ -84,6 +85,7 @@ def __init_finalize__(self, *args, function=None, **kwargs): # Data initialization initializer = kwargs.get('initializer') + if self.alias: self._initializer = None elif function is not None: @@ -91,6 +93,10 @@ def __init_finalize__(self, *args, function=None, **kwargs): # `f(x+1)`), so we just copy the reference to the original data self._initializer = None self._data = function._data + elif isinstance(self._allocator, DataReference): + # Don't want to reinitialise array if DataReference used as allocator; + # create a no-op intialiser to avoid overwriting the original array. + self._initializer = lambda x: None elif initializer is None or callable(initializer) or self.alias: # Initialization postponed until the first access to .data self._initializer = initializer diff --git a/tests/test_data.py b/tests/test_data.py index 3f4d0b5cb9..cf430dc744 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -6,9 +6,9 @@ switchconfig, SparseFunction, PrecomputedSparseFunction, PrecomputedSparseTimeFunction) from devito.data import LEFT, RIGHT, Decomposition, loc_data_idx, convert_index +from devito.data.allocators import DataReference from devito.tools import as_tuple from devito.types import Scalar -from devito.data.allocators import ExternalAllocator class TestDataBasic: @@ -1570,31 +1570,6 @@ def test_numpy_c_contiguous(): assert(u._data_allocated.flags.c_contiguous) -def test_external_allocator(): - shape = (2, 2) - space_order = 0 - numpy_array = np.ones(shape, dtype=np.float32) - g = Grid(shape) - f = Function(name='f', space_order=space_order, grid=g, - allocator=ExternalAllocator(numpy_array), initializer=lambda x: None) - - # Ensure the two arrays have the same value - assert(np.array_equal(f.data, numpy_array)) - - # Ensure the original numpy array is unchanged - assert(np.array_equal(numpy_array, np.ones(shape, dtype=np.float32))) - - # Change the underlying numpy array - numpy_array[:] = 3. - # Ensure the function.data changes too - assert(np.array_equal(f.data, numpy_array)) - - # Change the function.data - f.data[:] = 4. - # Ensure the underlying numpy array changes too - assert(np.array_equal(f.data, numpy_array)) - - def test_boolean_masking_array(): """ Test truth value of array, raised in Python 3.9 (MFE for issue #1788) @@ -1612,6 +1587,78 @@ def test_boolean_masking_array(): assert all(f.data == [1, 1, 0, 0, 1]) +class TestDataReference: + """ + Tests for passing data to a Function using a reference to a + preexisting array-like. + """ + + def test_w_array(self): + """Test using a preexisting NumPy array as Function data""" + grid = Grid(shape=(3, 3)) + a = np.reshape(np.arange(25, dtype=np.float32), (5, 5)) + b = a.copy() + c = a.copy() + + b[1:-1, 1:-1] += 1 + + f = Function(name='f', grid=grid, space_order=1, + allocator=DataReference(a)) + + # Check that the array hasn't been zeroed + assert np.any(a != 0) + + # Check that running operator updates the original array + Operator(Eq(f, f+1))() + assert np.all(a == b) + + # Check that updating the array updates the function data + a[1:-1, 1:-1] -= 1 + assert np.all(f.data_with_halo == c) + + def _w_data(self): + shape = (5, 5) + grid = Grid(shape=shape) + f = Function(name='f', grid=grid, space_order=1) + f.data_with_halo[:] = np.reshape(np.arange(49, dtype=np.float32), (7, 7)) + + g = Function(name='g', grid=grid, space_order=1, + allocator=DataReference(f._data)) + + # Check that the array hasn't been zeroed + assert np.any(f.data_with_halo != 0) + + assert np.all(f.data_with_halo == g.data_with_halo) + + # Update f + Operator(Eq(f, f+1))() + assert np.all(f.data_with_halo == g.data_with_halo) + + # Update g + Operator(Eq(g, g+1))() + assert np.all(f.data_with_halo == g.data_with_halo) + + check = np.array(f.data_with_halo[1:-1, 1:-1]) + + # Update both + Operator([Eq(f, f+1), Eq(g, g+1)])() + assert np.all(f.data_with_halo == g.data_with_halo) + # Check that it was incremented by two + check += 2 + assert np.all(f.data == check) + + def test_w_data(self): + """Test passing preexisting Function data to another Function""" + self._w_data() + + @pytest.mark.parallel(mode=[2, 4]) + def test_w_data_mpi(self, mode): + """ + Test passing preexisting Function data to another Function with MPI. + """ + self._w_data() + + if __name__ == "__main__": configuration['mpi'] = True TestDataDistributed().test_misc_data() diff --git a/tests/test_docstrings.py b/tests/test_docstrings.py index f4a34040f2..e95e786ace 100644 --- a/tests/test_docstrings.py +++ b/tests/test_docstrings.py @@ -19,7 +19,7 @@ @pytest.mark.parametrize('modname', [ 'types.basic', 'types.dimension', 'types.constant', 'types.grid', 'types.dense', 'types.sparse', 'types.equation', 'types.relational', 'operator', - 'data.decomposition', 'finite_differences.finite_difference', + 'data.decomposition', 'data.allocators', 'finite_differences.finite_difference', 'finite_differences.coefficients', 'finite_differences.derivative', 'ir.support.space', 'data.utils', 'data.allocators', 'builtins', 'symbolics.inspection', 'tools.utils', 'tools.data_structures' diff --git a/tests/test_rebuild.py b/tests/test_rebuild.py new file mode 100644 index 0000000000..0b29ea7047 --- /dev/null +++ b/tests/test_rebuild.py @@ -0,0 +1,42 @@ +import numpy as np + +from devito import Dimension, Function +from devito.data.allocators import DataReference + + +class TestFunction: + """Tests for rebuilding of Function types.""" + + def test_w_new_dims(self): + x = Dimension('x') + y = Dimension('y') + x0 = Dimension('x0') + y0 = Dimension('y0') + + f = Function(name='f', dimensions=(x, y), shape=(11, 11)) + f.data[:] = 1 + + dims0 = (x0, y0) + dims1 = (x, y0) + + f0 = f._rebuild(dimensions=dims0) + f1 = f._rebuild(dimensions=dims1) + f2 = f._rebuild(dimensions=f.dimensions) + f3 = f._rebuild(dimensions=dims0, + allocator=DataReference(f._data)) + + assert f0.function is f0 + assert f0.dimensions == dims0 + assert np.all(f0.data[:] == 0) + + assert f1.function is f1 + assert f1.dimensions == dims1 + assert np.all(f1.data[:] == 0) + + assert f2.function is f + assert f2.dimensions == f.dimensions + assert np.all(f2.data[:] == 1) + + assert f3.function is f3 + assert f3.dimensions == dims0 + assert np.all(f3.data[:] == 1)