data: Tie the lifetime of C-allocated memory to its views

devito/data/allocators.py

@@ -53,7 +53,9 @@ def alloc(self, shape, dtype, padding=0):
         ndarray, memfree_args
             The first element of the tuple is a numpy array that uses the
             allocated memory underneath. The second element is an opaque
-            object that is needed only for the "memfree" call.
+            object that is needed only for the "memfree" call; it must not
+            refer to the array itself, whose collection is what triggers the
+            "memfree" call in the first place.
         """
         return
 

devito/data/data.py

@@ -1,3 +1,4 @@
+import weakref
 from collections.abc import Iterable
 
 import numpy as np
@@ -72,21 +73,16 @@ def __new__(cls, shape, dtype, decomposition=None, modulo=None,
         assert all(i is None for i, j in zip(obj._decomposition, obj._modulo, strict=True)
                    if j is True)
 
-        return obj
+        if memfree_args is not None:
+            # Release the memory once `ndarray` is gone, and with it every view
+            # of it -- this Data, but also whatever was handed out by e.g.
+            # `Function.data` or `Function._data_allocated`, since NumPy anchors
+            # the `base` chain of all of them on `ndarray`. Releasing it when
+            # this Data dies instead would leave those views pointing at freed
+            # memory as soon as the owning Function went out of scope
+            weakref.finalize(ndarray, allocator.free, *memfree_args)
 
-    def __del__(self):
-        if getattr(self, "_memfree_args", None) is None:
-            # NOTE: The need for `getattr`, in place of `self._memfree_args`, was
-            # suggested for the first time in issue #1184. However, it appears
-            # that even though, as described in the issue, we initialize the
-            # attribute in `__array_finalize__`, an AttributeError exception may
-            # still be raised in some obscure situations. Our best explanation
-            # so far is that this is due to (un)pickling (as often used in a
-            # Dask/Distributed context), which may (re)create a Data object
-            # without going through `__array_finalize__`
-            return
-        self._allocator.free(*self._memfree_args)
-        self._memfree_args = None
+        return obj
 
     def __reduce__(self):
         warning("Pickling of `Data` objects is not supported. Casting to `numpy.ndarray`")
@@ -103,8 +99,8 @@ def __array_finalize__(self, obj):
         self._index_stash = None
 
         # Views or references created via operations on `obj` do not get an
-        # explicit reference to the underlying data (`_memfree_args`). This makes sure
-        # that only one object (the "root" Data) will free the C-allocated memory
+        # explicit reference to the underlying allocation (`_memfree_args`);
+        # only the "root" Data carries it
         self._memfree_args = None
 
         if not issubclass(type(obj), Data):

tests/test_data.py

@@ -1,13 +1,15 @@
+import gc
+
 import numpy as np
 import pytest
 
 from devito import (  # noqa
     ALLOC_ALIGNED, ALLOC_GUARD, Dimension, Eq, Function, Grid, Operator,
     PrecomputedSparseFunction, PrecomputedSparseTimeFunction, SparseFunction,
-    SparseTimeFunction, TimeFunction, configuration, switchconfig
+    SparseTimeFunction, TimeFunction, clear_cache, configuration, switchconfig
 )
 from devito.data import LEFT, RIGHT, Decomposition, convert_index, loc_data_idx
-from devito.data.allocators import DataReference
+from devito.data.allocators import DataReference, PosixAllocator
 from devito.data.distributed.layout import Layout
 from devito.data.distributed.selection import (
     Affine, Explicit, IndexScalar, Selection, index_has_array, result_dims
@@ -1965,6 +1967,123 @@ def test_scalar_arg_substitution():
     assert t0.subs('t0', t1) == t1
 
 
+class TestMemoryLifetime:
+
+    """
+    The memory backing a Function must outlive every view handed out of it.
+    """
+
+    class TrackingAllocator(PosixAllocator):
+
+        """A PosixAllocator that counts how many times it releases memory."""
+
+        def __init__(self):
+            super().__init__()
+            self.nfree = 0
+
+        def free(self, *args):
+            self.nfree += 1
+            super().free(*args)
+
+    def test_view_not_clobbered_by_later_allocations(self):
+        """
+        Hold on to a view, drop the Function, then allocate. With nothing
+        owning the memory, the view ends up reading whatever landed on the
+        freed block.
+        """
+        grid = Grid(shape=(12, 11, 10), dtype=np.float32)
+
+        f = Function(name='f', grid=grid, space_order=2)
+        f.data[:] = 1.
+
+        view = f._data_allocated
+        expected = view.copy()
+
+        del f
+        clear_cache()
+        gc.collect()
+
+        # Anything allocating the same shape lands on the freed block
+        others = [Function(name=f'g{i}', grid=grid, space_order=2) for i in range(8)]
+        for i in others:
+            i.data[:] = 9.
+
+        assert np.array_equal(view, expected)
+
+    def test_view_outlives_function(self):
+        """
+        A view handed out by a Function -- `data`, `data_with_halo`,
+        `_data_allocated` -- keeps the underlying memory alive, so that it is
+        safe to hold onto it after the Function itself is gone.
+        """
+        allocator = self.TrackingAllocator()
+
+        grid = Grid(shape=(4, 4))
+        f = Function(name='f', grid=grid, space_order=0, allocator=allocator)
+        f.data[:] = 3.
+
+        view = f._data_allocated
+
+        del f
+        clear_cache()
+        gc.collect()
+
+        assert allocator.nfree == 0
+        assert np.all(view == 3.)
+
+    def test_self_built_allocator(self):
+        """
+        Allocators that build the array themselves, by overriding `alloc`, are
+        covered by the same mechanism.
+        """
+        class SelfBuilding(self.TrackingAllocator):
+            def alloc(self, shape, dtype, padding=0):
+                # The free args must not refer to the array itself, or nothing
+                # would ever be collected
+                return (np.zeros(shape, dtype=dtype), (object(),))
+
+            def free(self, token):
+                self.nfree += 1
+
+        allocator = SelfBuilding()
+
+        grid = Grid(shape=(4, 4))
+        f = Function(name='f', grid=grid, space_order=0, allocator=allocator)
+        f.data[:] = 3.
+
+        view = f._data_allocated
+
+        del f
+        clear_cache()
+        gc.collect()
+        assert allocator.nfree == 0
+        assert np.all(view == 3.)
+
+        del view
+        gc.collect()
+        assert allocator.nfree == 1
+
+    def test_memory_released_with_last_view(self):
+        """
+        ... and the memory is released, exactly once, as soon as the last view
+        of it is gone.
+        """
+        allocator = self.TrackingAllocator()
+
+        grid = Grid(shape=(4, 4))
+        f = Function(name='f', grid=grid, space_order=0, allocator=allocator)
+        views = [f.data, f.data_with_halo, f._data_allocated]
+
+        del f
+        clear_cache()
+        gc.collect()
+        assert allocator.nfree == 0
+
+        del views
+        gc.collect()
+        assert allocator.nfree == 1
+
+
 @pytest.mark.skip(reason="will corrupt memory and risk crash")
 def test_oob_noguard():
     """