"""Memory estimates, and checks against the memory available to this process.
Call sites name what they allocate, for example the eigensystem of a matrix of
dimension ``D``, and the estimators below turn that into bytes. Each backend
adds the cost of its own lowering and solvers. :func:`require_memory` refuses
an allocation that cannot fit before quchip makes it.
"""
from __future__ import annotations
from math import prod
from pathlib import Path
from typing import Any
import numpy as np
import psutil
from scipy import sparse
from scipy.sparse.csgraph import reverse_cuthill_mckee
# Bytes of one complex double and of one 64-bit index.
COMPLEX_BYTES = 16
INDEX_BYTES = 8
# A dense Hermitian eigensystem keeps the matrix, a working copy that becomes
# the eigenvectors, the returned eigenvectors, and the complex and real
# workspace of LAPACK's divide-and-conquer driver.
_EIGENSYSTEM_COPIES = 6
# The cgroup v2 files of this process's container, when it has one.
_CGROUP = Path("/sys/fs/cgroup")
# SuperLU keeps a complex value and a 32-bit row index for each entry of its factors.
_LU_ENTRY_BYTES = COMPLEX_BYTES + 4
# SuperLU's fill depends on its column ordering and on partial pivoting. On
# quchip's stationary generators, from driven feedlines of transmons and
# resonators to a resonator with 300 levels, the fill was between 0.32 and 2.2
# times the reverse Cuthill-McKee envelope of the symmetrized pattern, so a
# quarter of the envelope stayed below each measured fill.
_ENVELOPE_FILL_FRACTION = 0.25
[docs]
def dense_bytes(*shape: int, copies: int = 1, itemsize: int = COMPLEX_BYTES) -> int:
"""Return the bytes of ``copies`` dense arrays of ``shape``.
Parameters
----------
*shape : int
Shape of one array.
copies : int, default=1
Number of arrays that exist at the same time.
itemsize : int, default=16
Bytes of one entry. The default is a complex double.
"""
return int(copies * itemsize * prod(shape))
[docs]
def eigensystem_bytes(dimension: int) -> int:
"""Return the peak bytes of the dense Hermitian eigensystem of a square matrix.
Parameters
----------
dimension : int
Number of rows of the matrix.
"""
return dense_bytes(dimension, dimension, copies=_EIGENSYSTEM_COPIES)
[docs]
def level_table_bytes(states: int, devices: int) -> int:
"""Return the bytes of a table that gives the level of each device in each product state.
Parameters
----------
states : int
Number of product states.
devices : int
Number of devices.
"""
return dense_bytes(states, devices, itemsize=INDEX_BYTES)
[docs]
def krylov_bytes(size: int, vectors: int) -> int:
"""Return the bytes of a Krylov basis.
Parameters
----------
size : int
Number of entries of one vector.
vectors : int
Number of vectors that the basis keeps.
"""
return dense_bytes(size, vectors)
[docs]
def sparse_lu_bytes(matrix: Any) -> int:
"""Return a low estimate of the bytes of the sparse LU factors of a square matrix.
SuperLU knows its fill only after the factorization, so the estimate is a
quarter of the reverse Cuthill-McKee envelope of the symmetrized pattern.
For a few modes with many levels the envelope grows faster than the fill,
and the estimate can exceed the fill of a resonator with many hundreds of
levels.
Parameters
----------
matrix : scipy.sparse matrix or array
Square matrix.
"""
size = matrix.shape[0]
if size == 0:
return 0
if matrix.format not in ("csr", "csc"):
matrix = sparse.csr_matrix(matrix)
# Read as CSR, a CSC matrix is its transpose, which has the same symmetrized pattern.
pattern = sparse.csr_matrix((np.ones(matrix.nnz, dtype=bool), matrix.indices, matrix.indptr), shape=(size, size))
pattern = (pattern + pattern.T).tocsr()
position = np.empty(size, dtype=np.int64)
position[reverse_cuthill_mckee(pattern, symmetric_mode=True)] = np.arange(size)
# After the reordering, the envelope of each row runs from its first entry to the diagonal.
first = position.copy()
rows = np.flatnonzero(np.diff(pattern.indptr))
if rows.size:
first[rows] = np.minimum(first[rows], np.minimum.reduceat(position[pattern.indices], pattern.indptr[rows]))
envelope = size + 2 * int(np.sum(position - first))
return int(_ENVELOPE_FILL_FRACTION * _LU_ENTRY_BYTES * envelope)
[docs]
def available_memory_bytes() -> int | None:
"""Return the memory this process can still allocate, or ``None`` when unknown.
The value is the available system memory that ``psutil`` reports. Under a
cgroup v2 memory limit, as in many Linux containers, the value is the
remaining limit when that is smaller.
"""
try:
limits = [int(psutil.virtual_memory().available)]
except (OSError, psutil.Error):
return None
try:
maximum = (_CGROUP / "memory.max").read_text().strip()
if maximum != "max":
limits.append(int(maximum) - int((_CGROUP / "memory.current").read_text()))
except (OSError, ValueError):
pass
return max(min(limits), 0)
[docs]
def fits_memory(required: float) -> bool:
"""Return whether an estimated allocation fits, and ``True`` when memory is unknown.
Parameters
----------
required : float
Estimated peak allocation in bytes.
"""
available = available_memory_bytes()
return available is None or required <= available
[docs]
def require_memory(required: float, *, task: str, remedy: str) -> None:
"""Raise :class:`MemoryError` before an allocation that cannot fit.
Parameters
----------
required : float
Estimated peak allocation in bytes.
task : str
What would allocate it, for the error message.
remedy : str
How to avoid the allocation.
"""
available = available_memory_bytes()
if available is not None and required > available:
raise MemoryError(
f"{task} needs about {required / 1e9:.1f} GB, but only {available / 1e9:.1f} GB "
f"is available. {remedy}"
)
