diff --git a/feectools/linalg/block.py b/feectools/linalg/block.py index 5352f93e9..f4085edc2 100644 --- a/feectools/linalg/block.py +++ b/feectools/linalg/block.py @@ -886,6 +886,14 @@ def nonzero_block_indices(self): """ return tuple(self._blocks) + # ... + @property + def nbytes(self): + """Local (per-MPI-rank) memory footprint of all non-zero blocks, in bytes. + Blocks which do not expose an 'nbytes' attribute (e.g. matrix-free operators) + are counted as zero.""" + return int(sum(getattr(Lij, 'nbytes', 0) for Lij in self._blocks.values())) + # ... def update_ghost_regions(self): for Lij in self._blocks.values(): diff --git a/feectools/linalg/kron.py b/feectools/linalg/kron.py index 6ee1ee4f5..80b013f14 100644 --- a/feectools/linalg/kron.py +++ b/feectools/linalg/kron.py @@ -73,6 +73,12 @@ def ndim( self ): def mats( self ): return self._mats + # ... + @property + def nbytes( self ): + """Local (per-MPI-rank) memory footprint of the 1d factor matrices, in bytes.""" + return int(sum(getattr(mat, 'nbytes', 0) for mat in self._mats)) + # ... def dot(self, x, out=None): diff --git a/feectools/linalg/memory.py b/feectools/linalg/memory.py new file mode 100644 index 000000000..2fc8982e9 --- /dev/null +++ b/feectools/linalg/memory.py @@ -0,0 +1,41 @@ +# coding: utf-8 +""" +Bookkeeping of the memory occupied by the stencil matrices that are currently alive. + +Every :class:`~feectools.linalg.stencil.StencilMatrix` that allocates its data array registers +itself (weakly) in the module-level :data:`stencil_matrix_memory` tracker, so that an application +can report how much memory its matrices actually take, without having to walk its own data +structures. Matrices created with ``dry_run=True`` do not allocate anything and are not registered. +""" + +import weakref + +__all__ = ('MatrixMemoryTracker', 'stencil_matrix_memory') + + +class MatrixMemoryTracker: + """Weak registry of allocated matrices; matrices that are garbage collected drop out of it.""" + + def __init__(self): + self._matrices = weakref.WeakSet() + + def register(self, matrix): + """Add a matrix to the registry (does not keep it alive).""" + self._matrices.add(matrix) + + def clear(self): + """Forget all registered matrices.""" + self._matrices.clear() + + @property + def n_matrices(self): + """Number of currently alive registered matrices.""" + return len(self._matrices) + + @property + def nbytes(self): + """Local (per-MPI-rank) memory footprint, in bytes, of all currently alive registered matrices.""" + return int(sum(matrix.nbytes for matrix in self._matrices)) + + +stencil_matrix_memory = MatrixMemoryTracker() diff --git a/feectools/linalg/stencil.py b/feectools/linalg/stencil.py index e13ca345d..6e3fe0bca 100644 --- a/feectools/linalg/stencil.py +++ b/feectools/linalg/stencil.py @@ -14,6 +14,7 @@ from feectools.ddm.mpi import mpi as MPI from feectools.linalg.basic import VectorSpace, Vector, LinearOperator +from feectools.linalg.memory import stencil_matrix_memory from feectools.ddm.cart import find_mpi_type, CartDecomposition, InterfaceCartDecomposition from feectools.ddm.utilities import get_data_exchanger from feectools.api.settings import PSYDAC_BACKENDS @@ -905,8 +906,14 @@ class StencilMatrix(LinearOperator): precompiled : bool Whether to use precompiled kernels for .dot() and .transpose() + + dry_run : bool + If True, only compute the shape of the data array (:attr:`data_shape`) and return early, + without allocating the (potentially large) data array and without setting up the + dot/transpose kernels. The resulting object is *not* usable as a linear operator; its only + purpose is to report the memory footprint the matrix would have via :attr:`nbytes`. """ - def __init__( self, V, W, pads=None , backend=None, precompiled=True): + def __init__( self, V, W, pads=None , backend=None, precompiled=True, dry_run=False): assert isinstance(V, StencilVectorSpace) assert isinstance(W, StencilVectorSpace) @@ -918,13 +925,21 @@ def __init__( self, V, W, pads=None , backend=None, precompiled=True): for p,vp in zip(pads, V.pads): assert p<=vp - self._pads = pads or tuple(V.pads) - dims = list(W.shape) - diags = [compute_diag_len(p, md, mc) for p,md,mc in zip(self._pads, V.shifts, W.shifts)] - self._data = xp.zeros(tuple(int(d) for d in (dims + diags)), dtype=W.dtype) - self._domain = V - self._codomain = W - self._ndim = len(dims) + self._pads = pads or tuple(V.pads) + dims = list(W.shape) + diags = [compute_diag_len(p, md, mc) for p,md,mc in zip(self._pads, V.shifts, W.shifts)] + self._data_shape = tuple(int(d) for d in (dims + diags)) + self._domain = V + self._codomain = W + self._ndim = len(dims) + self._dry_run = dry_run + + # memory estimation only: do not allocate the data array, see the nbytes property + if dry_run: + return + + self._data = xp.zeros(self._data_shape, dtype=W.dtype) + stencil_matrix_memory.register(self) self._backend = backend self._precompiled = precompiled self._is_T = False @@ -975,6 +990,17 @@ def __init__( self, V, W, pads=None , backend=None, precompiled=True): backend = PSYDAC_BACKENDS.get(os.environ.get('PSYDAC_BACKEND')) or PSYDAC_BACKENDS['python'] self.set_backend(backend, precompiled) + # ... + def __getattr__(self, name): + # only called when the attribute was not found the usual way; give a helpful + # message for the attributes that are missing on a dry-run matrix + if self.__dict__.get('_dry_run', False): + raise AttributeError( + f"'{type(self).__name__}.{name}' is not available because the matrix was created with " + "dry_run=True (memory estimation only, no data allocated)." + ) + raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'") + #-------------------------------------- # Abstract interface #-------------------------------------- @@ -992,6 +1018,28 @@ def codomain(self): def dtype(self): return self._domain.dtype + # ... + @property + def dry_run(self): + """Whether the matrix was created for memory estimation only, i.e. without allocating data.""" + return self._dry_run + + # ... + @property + def data_shape(self): + """Shape of the local data array (n_rows in each direction + n_diagonals in each direction).""" + return self._data_shape + + # ... + @property + def nbytes(self): + """Local (per-MPI-rank) memory footprint of the data array, in bytes. Also available + for matrices created with ``dry_run=True``, i.e. before/without allocating the data.""" + nbytes = xp.dtype(self._codomain.dtype).itemsize + for n in self._data_shape: + nbytes *= n + return int(nbytes) + # ... def dot(self, v, out=None): """ @@ -2076,6 +2124,11 @@ def codomain(self): def dtype(self): return self._data.dtype + @property + def nbytes(self): + """Local (per-MPI-rank) memory footprint of the data array, in bytes.""" + return int(self._data.nbytes) + def tosparse(self): return sp_diags(self._data.ravel()) @@ -2369,6 +2422,12 @@ def codomain(self): def dtype(self): return self.domain.dtype + # ... + @property + def nbytes(self): + """Local (per-MPI-rank) memory footprint of the data array, in bytes.""" + return int(self._data.nbytes) + # ... def dot(self, v, out=None): diff --git a/pyproject.toml b/pyproject.toml index a009e11b6..71cfd2299 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "feectools" -version = "0.1.6" +version = "0.1.7" description = "Slimmed-down fork of Psydac (https://github.com/pyccel/psydac) with less functionality and fewer dependencies." readme = "README.md" requires-python = ">= 3.10"