Merge pull request #2438 from devitocodes/func_rebuild

dsl: Improve Function rebuilding and use of preexisting arrays as Function data

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)

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

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,13 +85,18 @@ 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:
             # An object derived from user-level AbstractFunction (e.g.,
             # `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

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()

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'

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)