From 960c59eb9506c16c50d52a49e3792384103fde76 Mon Sep 17 00:00:00 2001 From: Mathias Louboutin Date: Sat, 5 Sep 2026 05:04:18 +0100 Subject: [PATCH 1/2] data: Tie the lifetime of C-allocated memory to its views `Data.__del__` released the memory as soon as the `Data` object died, but the views handed out of it outlive that. `Function._data_allocated` in particular returns `np.asarray(self._data)`, whose `base` NumPy collapses past the `Data` down to the array `alloc` built, so it keeps the `Data` alive not at all; holding onto one of those past the lifetime of the owning Function read freed memory. That array is, however, exactly what every view of the allocation -- the `Data` included -- does anchor on. So hang the deallocation off it with a `weakref.finalize`, and the memory outlives every view of it. --- devito/data/allocators.py | 4 +++- devito/data/data.py | 28 ++++++++++++---------------- 2 files changed, 15 insertions(+), 17 deletions(-) diff --git a/devito/data/allocators.py b/devito/data/allocators.py index 95fd8a0cd99..d883118f461 100644 --- a/devito/data/allocators.py +++ b/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 diff --git a/devito/data/data.py b/devito/data/data.py index 93963f33972..b41fc0eed3a 100644 --- a/devito/data/data.py +++ b/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): From 954a1fd41b821ed541c409ac8f69b5f0cd6644d9 Mon Sep 17 00:00:00 2001 From: Mathias Louboutin Date: Sat, 5 Sep 2026 05:04:18 +0100 Subject: [PATCH 2/2] tests: Cover the lifetime of a Function's data views --- tests/test_data.py | 123 ++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 121 insertions(+), 2 deletions(-) diff --git a/tests/test_data.py b/tests/test_data.py index 64069bf65d3..13f2ae5645a 100644 --- a/tests/test_data.py +++ b/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(): """