diff --git a/AGENTS.md b/AGENTS.md index b0c6d35..6af3405 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -20,9 +20,12 @@ tensor, so there is no wrapper type. `A`); that licence does not extend to anything else. - Comment invariants, contracts and non-obvious numerical choices only. Never narrate the code. -- Backend-agnostic code: go through `src_method.utils._backend` instead of - importing `numpy` or `cupy` directly in the algorithms, so CPU and GPU paths - stay in sync. +- Backend-agnostic code: use the array API standard, with the namespace taken + from the arrays (`array_api_compat.array_namespace`) or from the + `DeviceGroup`, never an `xp` parameter; annotate arrays as + `src_method.utils.Array`. What the standard lacks (streams, pinned memory, + CuPy's pool) goes through `src_method.utils._backend`. Never import `cupy` + in the algorithms, so CPU and GPU paths stay in sync. - Log with the standard library: `logger = logging.getLogger(__name__)` per module, lazy `%`-style arguments, never `print`. Never add handlers, levels or `basicConfig` in the package; progress messages go at `DEBUG`. diff --git a/README.md b/README.md index 82c2800..9990826 100644 --- a/README.md +++ b/README.md @@ -2,6 +2,8 @@ [![Documentation](https://github.com/Algorithmiq/src-method/actions/workflows/docpages.yml/badge.svg)](https://algorithmiq.github.io/src-method/) [![Test src_method](https://github.com/Algorithmiq/src-method/actions/workflows/test.yml/badge.svg)](https://github.com/Algorithmiq/src-method/actions/workflows/test.yml) +[![Coverage](https://sonarcloud.io/api/project_badges/measure?project=Algorithmiq_src-method&metric=coverage&token=ab544eca4948ef2f80dffa566c31a9c00026a825)](https://sonarcloud.io/summary/new_code?id=Algorithmiq_src-method) +[![Quality gate status](https://sonarcloud.io/api/project_badges/measure?project=Algorithmiq_src-method&metric=alert_status&token=ab544eca4948ef2f80dffa566c31a9c00026a825)](https://sonarcloud.io/summary/new_code?id=Algorithmiq_src-method) [![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](LICENSE) An implementation of the SRC algorithm introduced in [Camaño, Epperly and Tropp, *Quantum* **10**, 2022 (2026)](https://doi.org/10.22331/q-2026-03-10-2022), extending the idea to other kinds of tensor networks. @@ -109,21 +111,6 @@ On a node with several GPUs, `Resources(devices=4)` splits the sweep among four See [Large problems](https://algorithmiq.github.io/src-method/features/large-problems) for the budgets, the scratch directory, several GPUs and how to read the logged plan. -### Logging - -`src_method` logs through the standard library `logging` module, under the -`src_method` logger, and stays silent by default: it adds only a `NullHandler` -and never touches the root logger or its handlers. Progress and timing messages -are emitted at `DEBUG`, the small-network fallback at `WARNING`. To see them, -configure logging in your application: - -```python -import logging - -logging.basicConfig(level=logging.INFO) -logging.getLogger("src_method").setLevel(logging.DEBUG) -``` - ## Installation ```bash @@ -150,45 +137,9 @@ dependencies = ["src_method[gpu-nvidia]"] dependencies = ["src_method[gpu-rocm]"] ``` -## Setting up the development environment - -The code has a [DevContainer] configuration that will get you up and running -with all dependencies installed and configured, including sane defaults for the -editor. - -You will need: - -1. A working [Docker] installation: - - For macOS and Windows, install [Docker Desktop](https://docs.docker.com/get-docker/) - - For Linux, install [Docker Engine](https://docs.docker.com/engine/install/#server) following the instructions for your specific distro. -2. The [Visual Studio Code] editor. A recent version is recommended, _e.g._ >=1.78 -3. The VSCode [DevContainers extension](https://marketplace.visualstudio.com/items?itemName=ms-vscode-remote.remote-containers). -4. The [GitHub CLI] tool. - -You can clone the repository with: - -``` -gh repo clone Algorithmiq/src-method -``` - -We recommend using a Git credential manager, such as [GitHub CLI], configured to -use HTTPS as protocol for Git operations. - -Once the code is locally available, you can open its containing folder in -[Visual Studio Code]. The editor will then set up the [DevContainer] for you. -The first time you open the folder the startup will take a few minutes. Once the -process is done, you will have _all_ project dependencies installed, including -the git hooks. -[Visual Studio Code] will be already configured with all the extensions helpful for Python development. - -> [!TIP] -> The order in which Visual Studio Code loads the extensions in the -> DevContainer is non-deterministic. You might have to execute the *Reload Window* -> command to get everything to work as expected after a fresh build of the container. - -### Alternative: Nix flake +## Nix flake -If you prefer [Nix] over Docker, the repository ships a `flake.nix` that provides +The repository ships a `flake.nix` that provides a development shell with [uv], [Git] and the [GitHub CLI], plus the native libraries the binary wheels need at runtime. Python itself and all project dependencies remain managed by `uv`. @@ -219,28 +170,9 @@ The documentation is a [Fumadocs] site under `docs/`, published [at this link]. It covers concepts, features, tutorials, the API reference generated from the docstrings, and the developer guide. -To preview it locally you need [Node.js] 22 or later: - -```bash -uv sync --group docs -cd docs -npm ci -uv run python scripts/gen_api_dump.py src_method -d . -node scripts/generate-api.mjs -uv run python scripts/notebooks_to_mdx.py -npm run dev -``` - -and open http://localhost:3000. See [`docs/README.md`](docs/README.md) for details. - -[DevContainer]: https://containers.dev/ -[Docker]: https://docs.docker.com/get-docker/ -[Visual Studio Code]: https://code.visualstudio.com/ [GitHub CLI]: https://cli.github.com/ [Git]: https://git-scm.com/ -[Nix]: https://nixos.org/download/ [direnv]: https://direnv.net/ [uv]: https://docs.astral.sh/uv/ [Fumadocs]: https://fumadocs.dev -[Node.js]: https://nodejs.org [at this link]: https://algorithmiq.github.io/src-method/ diff --git a/docs/content/docs/contributing/dependencies.mdx b/docs/content/docs/contributing/dependencies.mdx index dea8cb9..220c7ca 100644 --- a/docs/content/docs/contributing/dependencies.mdx +++ b/docs/content/docs/contributing/dependencies.mdx @@ -7,8 +7,8 @@ Dependencies are managed with [uv](https://docs.astral.sh/uv/) and declared in `pyproject.toml`: - **Runtime dependencies** (`[project] dependencies`) are what users install: - NumPy and `opt_einsum`. Keep this list short; every entry is imposed on every - downstream project. + NumPy, `opt_einsum` and `array-api-compat`. Keep this list short; every entry is + imposed on every downstream project. - **Extras** (`[project.optional-dependencies]`) enable optional features in production: `gpu-nvidia` and `gpu-rocm` add CuPy. - **Dependency groups** (`[dependency-groups]`) are for development only and are diff --git a/docs/content/docs/contributing/index.mdx b/docs/content/docs/contributing/index.mdx index 3fa6b38..87318af 100644 --- a/docs/content/docs/contributing/index.mdx +++ b/docs/content/docs/contributing/index.mdx @@ -74,5 +74,9 @@ not just the shapes. The project targets Python 3.11+, is checked with `ruff` under a strict rule set and type-checked with [`ty`](https://github.com/astral-sh/ty). Public functions carry type hints and Google-style docstrings without types in the argument list. -The algorithms go through `src_method.utils._backend` rather than importing NumPy -or CuPy directly, so the CPU and GPU paths stay in sync. +The algorithms speak the [array API standard](https://data-apis.org/array-api/) +through `array-api-compat`: they take the namespace from their arrays +(`array_namespace(x)`) instead of importing NumPy or CuPy, and annotate them as +`src_method.utils.Array`. What the standard leaves out (streams, page-locked +memory, CuPy's memory pool, host-side random draws) lives in +`src_method.utils._backend`, so the CPU and GPU paths stay in sync. diff --git a/docs/content/docs/getting-started.mdx b/docs/content/docs/getting-started.mdx index f8770af..8c45bcd 100644 --- a/docs/content/docs/getting-started.mdx +++ b/docs/content/docs/getting-started.mdx @@ -16,8 +16,8 @@ uv pip install "src_method[gpu-nvidia]" uv pip install "src_method[gpu-rocm]" ``` -The runtime dependencies are NumPy and `opt_einsum`. The GPU extras add CuPy; -see [GPU execution](/features/gpu). +The runtime dependencies are NumPy, `opt_einsum` and `array-api-compat`. The GPU +extras add CuPy; see [GPU execution](/features/gpu). ## The three entry points diff --git a/pyproject.toml b/pyproject.toml index 1297b30..f4b1098 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,7 +32,7 @@ classifiers = [ "Typing :: Typed", ] dynamic = ["version"] -dependencies = ["numpy>=2.0,<3", "opt_einsum>=3.4.0"] +dependencies = ["array-api-compat>=1.12", "numpy>=2.0,<3", "opt_einsum>=3.4.0"] [project.optional-dependencies] diff --git a/src/src_method/_group.py b/src/src_method/_group.py index 8e9facb..20c6602 100644 --- a/src/src_method/_group.py +++ b/src/src_method/_group.py @@ -19,13 +19,21 @@ from concurrent.futures import ThreadPoolExecutor from typing import TYPE_CHECKING, Any, NamedTuple, TypeVar -from .utils import copy_into, current_stream, enable_peer_access, use_device +from ._plan import device_ids +from .utils import ( + copy_into, + current_device, + current_stream, + device_count, + enable_peer_access, + get_xp, + use_device, +) if TYPE_CHECKING: from collections.abc import Callable, Sequence - from types import ModuleType - from numpy.typing import NDArray + from .utils import Array T = TypeVar("T") @@ -60,17 +68,26 @@ class DeviceGroup: """The devices of one sweep and the collectives between them. Args: - xp: Array module (``numpy`` or ``cupy``). - devices: The device ids, in rank order; rank 0 is the root. + device: ``"cpu"`` or ``"gpu"``, as in `src_method.stack.src`. + devices: A count ``n`` for the first ``n`` visible devices, or the device + ids in rank order; rank 0 is the root. Defaults to the current device. + + Raises: + ValueError: If a device does not exist. """ - def __init__(self, xp: ModuleType, devices: Sequence[int]) -> None: - self.xp = xp - self.devices = tuple(devices) + def __init__(self, device: str, devices: int | Sequence[int] | None = None) -> None: + self.xp = get_xp(device) + ids = device_ids(devices) or (current_device(self.xp),) + count = device_count(self.xp) + if count is not None and max(ids) >= count: + msg = f"Requested devices {list(ids)}, but only {count} are visible." + raise ValueError(msg) + self.devices = ids self.size = len(self.devices) # Uploading 1/size of each core and gathering the rest pays off only when # the devices copy directly between them. - self.split_uploads = self.size > 1 and enable_peer_access(xp, self.devices) + self.split_uploads = self.size > 1 and enable_peer_access(self.xp, ids) self._barrier = threading.Barrier(self.size) # Two sets of slots, alternating between consecutive collectives: a rank may # publish into the next collective while slower ranks still read this one. @@ -171,12 +188,12 @@ def _settle(self, rank: int, parity: int) -> None: stream.wait_event(event) self._parts[parity][rank] = None - def _pull(self, part: _Part, dst: NDArray) -> None: + def _pull(self, part: _Part, dst: Array) -> None: stream = current_stream(self.xp) stream.wait_event(part.ready) - copy_into(dst, part.array, self.xp, stream) + copy_into(dst, part.array, stream) - def allgather(self, rank: int, local: NDArray) -> NDArray: + def allgather(self, rank: int, local: Array) -> Array: """Concatenate the blocks of every rank along the first axis, on every rank. Args: @@ -199,7 +216,7 @@ def allgather(self, rank: int, local: NDArray) -> NDArray: self._settle(rank, parity) return full - def gather(self, rank: int, local: NDArray, root: int = 0) -> NDArray | None: + def gather(self, rank: int, local: Array, root: int = 0) -> Array | None: """Concatenate the blocks of every rank along the first axis, on ``root``. Args: @@ -224,7 +241,7 @@ def gather(self, rank: int, local: NDArray, root: int = 0) -> NDArray | None: self._settle(rank, parity) return full - def scatter(self, rank: int, full: NDArray | None, root: int = 0) -> NDArray: + def scatter(self, rank: int, full: Array | None, root: int = 0) -> Array: """Give every rank its block (see `blocks`) of the first axis of ``full``. Args: @@ -246,7 +263,7 @@ def scatter(self, rank: int, full: NDArray | None, root: int = 0) -> NDArray: self._settle(rank, parity) return local - def _empty_like_parts(self, parts: Sequence[_Part]) -> NDArray: + def _empty_like_parts(self, parts: Sequence[_Part]) -> Array: first = parts[0].array rows = sum(part.array.shape[0] for part in parts) return self.xp.empty((rows, *first.shape[1:]), dtype=first.dtype) diff --git a/src/src_method/_kernels.py b/src/src_method/_kernels.py index 47190d4..85da865 100644 --- a/src/src_method/_kernels.py +++ b/src/src_method/_kernels.py @@ -20,9 +20,10 @@ if TYPE_CHECKING: from collections.abc import Sequence - from numpy.typing import NDArray from opt_einsum.contract import ContractExpression + from .utils import Array + Shape = tuple[int, ...] @@ -71,7 +72,7 @@ class _Contractions: def __init__(self) -> None: self._compiled: dict[tuple[str, tuple[Shape, ...]], ContractExpression] = {} - def __call__(self, eq: str, *operands: NDArray) -> NDArray: + def __call__(self, eq: str, *operands: Array) -> Array: shapes = tuple(op.shape for op in operands) expr = self._compiled.get((eq, shapes)) if expr is None: @@ -90,7 +91,7 @@ def __init__(self, depth: int) -> None: self.eqs = equations(depth) self._contract = _Contractions() - def env(self, env: NDArray, omega: NDArray, cores: Sequence[NDArray]) -> NDArray: + def env(self, env: Array, omega: Array, cores: Sequence[Array]) -> Array: """Advance a batch of sketch columns of the environment by one site. Args: @@ -103,7 +104,7 @@ def env(self, env: NDArray, omega: NDArray, cores: Sequence[NDArray]) -> NDArray """ return self._contract(self.eqs.ltr, env, omega, *cores) - def sketch(self, env: NDArray, cores: Sequence[NDArray], proj: NDArray) -> NDArray: + def sketch(self, env: Array, cores: Sequence[Array], proj: Array) -> Array: """Sketch a batch of columns of the site's running core. Args: @@ -116,7 +117,7 @@ def sketch(self, env: NDArray, cores: Sequence[NDArray], proj: NDArray) -> NDArr """ return self._contract(self.eqs.rtl_m, env, *cores, proj) - def project(self, eta: NDArray, cores: Sequence[NDArray], proj: NDArray) -> NDArray: + def project(self, eta: Array, cores: Sequence[Array], proj: Array) -> Array: """Project a batch of rows of the new projected environment. Args: @@ -129,7 +130,7 @@ def project(self, eta: NDArray, cores: Sequence[NDArray], proj: NDArray) -> NDAr """ return self._contract(self.eqs.rtl_s, eta, *cores, proj) - def first(self, cores: Sequence[NDArray], proj: NDArray) -> NDArray: + def first(self, cores: Sequence[Array], proj: Array) -> Array: """Contract the first site with a batch of rows of ``S``. Args: diff --git a/src/src_method/_plan.py b/src/src_method/_plan.py index 8a2a099..0104e35 100644 --- a/src/src_method/_plan.py +++ b/src/src_method/_plan.py @@ -23,19 +23,13 @@ import numpy as np from ._kernels import equations, peak_elements -from .utils import ( - current_device, - device_count, - device_memory, - host_memory_available, - is_host, - use_device, -) +from .utils import device_memory, host_memory_available, is_host, use_device if TYPE_CHECKING: import os from collections.abc import Callable - from types import ModuleType + + from ._group import DeviceGroup Tier = Literal["device", "host", "disk"] Shape = tuple[int, ...] @@ -137,10 +131,10 @@ def __post_init__(self) -> None: for value in (self.gpu_memory, self.host_memory): if value is not None: parse_size(value) - _device_ids(self.devices) + device_ids(self.devices) -def _device_ids(devices: object) -> tuple[int, ...] | None: +def device_ids(devices: object) -> tuple[int, ...] | None: """Validate a device request and return its ids, or ``None`` for the default.""" if devices is None: return None @@ -164,29 +158,6 @@ def _device_ids(devices: object) -> tuple[int, ...] | None: return tuple(ids) -def resolve_devices(resources: Resources | None, xp: ModuleType) -> tuple[int, ...]: - """Return the device ids of a sweep, in rank order. - - Args: - resources: The requested devices, or ``None`` for the current device. - xp: Array module (``numpy`` or ``cupy``). - - Returns: - The device ids. - - Raises: - ValueError: If a device does not exist. - """ - ids = _device_ids(resources.devices if resources is not None else None) - if ids is None: - return (current_device(xp),) - count = device_count(xp) - if count is not None and max(ids) >= count: - msg = f"Requested devices {list(ids)}, but only {count} are visible." - raise ValueError(msg) - return ids - - @dataclass(frozen=True) class Budgets: """The resolved budgets of one sweep, in bytes. @@ -210,15 +181,12 @@ class Budgets: device_cap: int | None = None -def resolve_budgets( - resources: Resources | None, xp: ModuleType, devices: Sequence[int] = (0,) -) -> Budgets: +def resolve_budgets(resources: Resources | None, group: DeviceGroup) -> Budgets: """Turn `Resources` into byte budgets, detecting those left unset. Args: resources: The requested budgets, or ``None`` to detect all of them. - xp: Array module (``numpy`` or ``cupy``). - devices: The devices of the sweep. The device budget is that of each of + group: The devices of the sweep. The device budget is that of each of them: the smallest detected one. Host memory and disk are those of the node; on the CPU the simulated devices split the host budget. @@ -235,15 +203,15 @@ def resolve_budgets( host = parse_size(resources.host_memory) else: host = int(host_memory_available() * (1 - HOST_MARGIN_FRACTION)) - unified = is_host(xp) + unified = is_host(group.xp) cap = None if unified: - device = max(host, 0) // len(devices) + device = max(host, 0) // group.size else: caps = [] - for d in devices: - with use_device(xp, d): - available, total = device_memory(xp) + for d in group.devices: + with use_device(group.xp, d): + available, total = device_memory(group.xp) margin = max(int(GPU_MARGIN_FRACTION * total), GPU_MARGIN_MIN) caps.append((available - margin, margin)) margin = min(m for _, m in caps) diff --git a/src/src_method/_sites.py b/src/src_method/_sites.py index 531be36..75f2f04 100644 --- a/src/src_method/_sites.py +++ b/src/src_method/_sites.py @@ -35,11 +35,10 @@ if TYPE_CHECKING: from collections.abc import Sequence - from types import ModuleType, TracebackType - - from numpy.typing import NDArray + from types import TracebackType from ._tensor_train import Site + from .utils import Array # Byte alignment of each core inside a staging buffer. _ALIGN = 256 @@ -123,26 +122,24 @@ class SiteSource: Args: layers: A stack in ket form, each layer a sequence of array-likes with ``shape``, ``dtype``, ``ndim`` and ``np.asarray`` support. - xp: Array module (``numpy`` or ``cupy``). order: The sites, in the order they will be requested. depth: How many sites the loader may read ahead of the one in use. - group: The devices that share the cores; a single device by default. + group: The devices that share the cores. """ def __init__( self, layers: Sequence[Sequence[Site]], - xp: ModuleType, order: Sequence[int], *, depth: int, - group: DeviceGroup | None = None, + group: DeviceGroup, ) -> None: self._layers = layers self._kinds = [known_kind(layer) for layer in layers] self._last = len(layers[0]) - 1 - self._xp = xp - self._group = group or DeviceGroup(xp, [0]) + self._xp = xp = group.xp + self._group = group self._order = list(order) self._staged: list[Future[_Staged]] = [Future() for _ in self._order] self._taken = [0] * len(self._order) @@ -259,7 +256,7 @@ def _release(self, k: int, event: Any) -> None: # noqa: ANN401 # --- reading --- - def next(self, rank: int = 0) -> tuple[NDArray, ...]: + def next(self, rank: int = 0) -> tuple[Array, ...]: """Return the padded cores of the next site in order, on the rank's device. Args: @@ -290,7 +287,7 @@ def next(self, rank: int = 0) -> tuple[NDArray, ...]: for core, kind in zip(cores, self._kinds) ) - def _upload(self, rank: int, k: int, staged: _Staged) -> list[NDArray]: + def _upload(self, rank: int, k: int, staged: _Staged) -> list[Array]: xp, stream = self._xp, current_stream(self._xp) if not self._group.split_uploads: cores = [ @@ -311,7 +308,7 @@ def _upload(self, rank: int, k: int, staged: _Staged) -> list[NDArray]: ] -def _read(core: Site) -> NDArray: +def _read(core: Site) -> np.ndarray: """Bring a core into memory; memmaps are read now, not on first touch.""" if isinstance(core, np.memmap) or not isinstance(core, np.ndarray): return np.array(core) diff --git a/src/src_method/_store.py b/src/src_method/_store.py index 1abbf3e..0f9fb5c 100644 --- a/src/src_method/_store.py +++ b/src/src_method/_store.py @@ -21,6 +21,7 @@ from typing import TYPE_CHECKING, Any, Self import numpy as np +from array_api_compat import array_namespace from .utils import current_stream, pinned_empty, to_device_async, to_host_async @@ -29,9 +30,8 @@ from pathlib import Path from types import ModuleType, TracebackType - from numpy.typing import NDArray - from ._plan import Plan + from .utils import Array # Page-locked buffers per direction: one in use, one in flight. RING_SIZE = 2 @@ -47,13 +47,13 @@ class EnvironmentStore: """Keep the environments ``C_j`` of a sweep, batch by batch. Use as a context manager: leaving it waits for pending writes, stops the - threads and removes the scratch files, also after an exception. + threads and removes the scratch files, also after an exception. The device + arrays it returns are of the backend of those it was given. Args: plan: The plan of the sweep, which fixes the tier and batches of each site. env_shapes: The shape ``(chi, *left_bonds)`` of ``C_j`` for every site. dtype: The data type of the environments. - xp: Array module (``numpy`` or ``cupy``). scratch_dir: Where the disk tier creates its per-process directory. copy_stream: The stream that performs the transfers. """ @@ -63,7 +63,6 @@ def __init__( plan: Plan, env_shapes: Sequence[tuple[int, ...]], dtype: type | np.dtype, - xp: ModuleType, scratch_dir: Path, *, copy_stream: Any, # noqa: ANN401 (a cupy or null stream) @@ -72,9 +71,10 @@ def __init__( self._tiers = [site.tier for site in plan.sites] self._shapes = [tuple(shape) for shape in env_shapes] self._dtype = np.dtype(dtype) - self._xp = xp + # The namespace of the staged arrays, known from the first staged put. + self._xp: ModuleType | None = None self._copy = copy_stream - self._device: dict[int, NDArray] = {} + self._device: dict[int, Array] = {} self._host: dict[int, np.ndarray] = {} self._fds: dict[int, int] = {} self._writes: dict[int, list[Future[None]]] = {} @@ -98,11 +98,10 @@ def __init__( self._reader = ThreadPoolExecutor(1, thread_name_prefix="src-env-reader") self._out: queue.Queue[Any] = queue.Queue() self._in: queue.Queue[tuple[Any, Any]] = queue.Queue() - if staged: - n_bytes = max(self._batch_rows(j) * self._row_bytes(j) for j in staged) - for _ in range(RING_SIZE): - self._out.put(pinned_empty(n_bytes, xp)) - self._in.put((pinned_empty(n_bytes, xp), None)) + # The page-locked rings are allocated by the first staged put. + self._ring_bytes = max( + (self._batch_rows(j) * self._row_bytes(j) for j in staged), default=0 + ) def _row_bytes(self, j: int) -> int: return prod(self._shapes[j][1:]) * self._dtype.itemsize @@ -149,7 +148,7 @@ def close(self, *, wait: bool = True) -> None: # --- writing --- - def put(self, j: int, lo: int, hi: int, x: NDArray) -> None: + def put(self, j: int, lo: int, hi: int, x: Array) -> None: """Store columns ``lo:hi`` of ``C_j``. Args: @@ -160,17 +159,24 @@ def put(self, j: int, lo: int, hi: int, x: NDArray) -> None: """ if self._tiers[j] == "device": if j not in self._device: - self._device[j] = self._xp.empty(self._shapes[j], dtype=self._dtype) + xp = array_namespace(x) + self._device[j] = xp.empty(self._shapes[j], dtype=self._dtype) self._device[j][lo:hi] = x return + xp = self._xp = array_namespace(x) self._raise_failed() self._open(j) + if self._ring_bytes: + for _ in range(RING_SIZE): + self._out.put(pinned_empty(self._ring_bytes, xp)) + self._in.put((pinned_empty(self._ring_bytes, xp), None)) + self._ring_bytes = 0 start = perf_counter() buffer = self._out.get() self.stall_seconds += perf_counter() - start host = buffer[: x.nbytes].view(self._dtype).reshape(x.shape) # The copy starts once the kernel that wrote x is done. - self._copy.wait_event(current_stream(self._xp).record()) + self._copy.wait_event(current_stream(xp).record()) to_host_async(x, host, self._copy) done = self._copy.record() # x stays referenced by the task until the copy is complete. @@ -195,7 +201,7 @@ def _write( self, j: int, lo: int, - x: NDArray, + x: Array, *, host: np.ndarray, done: Any, # noqa: ANN401 (a cupy or null event) @@ -283,7 +289,7 @@ def _read( ) return buffer, host - def get(self, j: int, lo: int, hi: int) -> NDArray: + def get(self, j: int, lo: int, hi: int) -> Array: """Return columns ``lo:hi`` of ``C_j`` on the device. Args: @@ -312,8 +318,10 @@ def get(self, j: int, lo: int, hi: int) -> NDArray: self.stall_seconds += perf_counter() - start # On the current stream, so that the copy is ordered with the kernels that # use it and with the ones that freed the memory it reuses. - stream = current_stream(self._xp) - device = to_device_async(host, self._xp, stream) + xp = self._xp + assert xp is not None # noqa: S101 (C_j was put before it is read) + stream = current_stream(xp) + device = to_device_async(host, xp, stream) self._in.put((buffer, stream.record())) self._schedule() return device diff --git a/src/src_method/_sweep.py b/src/src_method/_sweep.py index c023cd7..114f1f9 100644 --- a/src/src_method/_sweep.py +++ b/src/src_method/_sweep.py @@ -33,10 +33,11 @@ from typing import TYPE_CHECKING import numpy as np +from array_api_compat import array_namespace from ._group import DeviceGroup, blocks from ._kernels import SiteKernels -from ._plan import make_plan, resolve_budgets, resolve_devices +from ._plan import make_plan, resolve_budgets from ._sites import SiteSource, padded_shapes, site_bytes, sweep_order from ._store import EnvironmentStore from ._tensor_train import unpad @@ -52,12 +53,12 @@ if TYPE_CHECKING: from collections.abc import Callable, Iterator, Sequence - from types import ModuleType - from numpy.typing import DTypeLike, NDArray + from numpy.typing import DTypeLike from ._plan import Budgets, Plan, Resources from ._tensor_train import Site, TrainKind + from .utils import Array logger = logging.getLogger(__name__) @@ -77,53 +78,59 @@ class _SplittingKernels: Args: kernels: The kernels of the sweep. - xp: Array module (``numpy`` or ``cupy``). tag: Prefix of the log messages. """ - def __init__(self, kernels: SiteKernels, xp: ModuleType, tag: str) -> None: - self._kernels, self._xp, self._tag = kernels, xp, tag + def __init__(self, kernels: SiteKernels, tag: str) -> None: + self._kernels, self._tag = kernels, tag self._pieces: dict[str, int] = {} - def env(self, env: NDArray, omega: NDArray, cores: Sequence[NDArray]) -> NDArray: + def env(self, env: Array, omega: Array, cores: Sequence[Array]) -> Array: """`SiteKernels.env`, split along the sketch columns.""" return self._split( "env", - env.shape[0], + env, lambda a, b: self._kernels.env(env[a:b], omega[a:b], cores), axis=0, ) - def sketch(self, env: NDArray, cores: Sequence[NDArray], proj: NDArray) -> NDArray: + def sketch(self, env: Array, cores: Sequence[Array], proj: Array) -> Array: """`SiteKernels.sketch`, split along the sketch columns.""" return self._split( "sketch", - env.shape[0], + env, lambda a, b: self._kernels.sketch(env[a:b], cores, proj), axis=-1, ) - def project(self, eta: NDArray, cores: Sequence[NDArray], proj: NDArray) -> NDArray: + def project(self, eta: Array, cores: Sequence[Array], proj: Array) -> Array: """`SiteKernels.project`, split along the rows.""" return self._split( "project", - eta.shape[0], + eta, lambda a, b: self._kernels.project(eta[a:b], cores, proj), axis=0, ) - def first(self, cores: Sequence[NDArray], proj: NDArray) -> NDArray: + def first(self, cores: Sequence[Array], proj: Array) -> Array: """`SiteKernels.first`, split along the rows.""" return self._split( "first", - proj.shape[0], + proj, lambda a, b: self._kernels.first(cores, proj[a:b]), axis=-3, ) def _split( - self, kernel: str, n: int, compute: Callable[[int, int], NDArray], *, axis: int - ) -> NDArray: + self, + kernel: str, + batched: Array, + compute: Callable[[int, int], Array], + *, + axis: int, + ) -> Array: + """Compute ``compute`` over the first axis of ``batched``, in pieces.""" + n = batched.shape[0] parts, start = [], 0 while start < n: pieces = self._pieces.get(kernel, 1) @@ -133,7 +140,7 @@ def _split( except MemoryError: if stop - start == 1: raise - release_cached_blocks(self._xp) + release_cached_blocks(array_namespace(batched)) self._pieces[kernel] = 2 * pieces logger.debug( "%sOut of memory in a %s batch of %d rows; %d pieces from now on", @@ -144,7 +151,9 @@ def _split( ) continue start = stop - return parts[0] if len(parts) == 1 else self._xp.concatenate(parts, axis=axis) + if len(parts) == 1: + return parts[0] + return array_namespace(batched).concat(parts, axis=axis) def sweep( @@ -152,12 +161,12 @@ def sweep( kind: TrainKind, chi_out: int, prng: np.random.Generator, - xp: ModuleType, *, + device: str = "cpu", cutoff: float = 0.0, dtype: DTypeLike, resources: Resources | None = None, -) -> list[NDArray]: +) -> list[np.ndarray]: """Contract and compress a stack in ket form with one SRC sweep. Args: @@ -169,7 +178,7 @@ def sweep( chi_out: The sketch size, which is the maximum output bond dimension. prng: The generator for the Gaussian sketches, always host-side so that a seed gives the same draws on every device. - xp: Array module (``numpy`` or ``cupy``). + device: ``"cpu"`` or ``"gpu"``, as in `src_method.stack.src`. cutoff: Relative singular-value cutoff for adaptive bond truncation. dtype: The data type of the sketches. resources: Memory budgets, scratch space and devices; detected when @@ -183,9 +192,8 @@ def sweep( # The dtype of the environments and the output: the sketches promoted by the # cores, as the contractions would. work = np.result_type(dtype, *(site.dtype for layer in layers for site in layer)) - devices = resolve_devices(resources, xp) - group = DeviceGroup(xp, devices) - budgets = resolve_budgets(resources, xp, devices) + group = DeviceGroup(device, resources.devices if resources else None) + budgets = resolve_budgets(resources, group) plan = make_plan( shapes, site_bytes(layers), chi_out, work, budgets, devices=group.size ) @@ -194,7 +202,7 @@ def sweep( logger.debug( "SRC plan: devices=%s, prefetch=%s, device peak=%d B, host peak=%d B, " "disk=%d B, scratch=%s, tiers=%s, batches (env, sketch, project)=%s", - list(devices), + list(group.devices), plan.prefetch, plan.device_peak, plan.host_peak, @@ -218,7 +226,7 @@ def sweep( depth=len(layers), ) with SiteSource( - layers, xp, sweep_order(len(shapes)), depth=plan.prefetch, group=group + layers, sweep_order(len(shapes)), depth=plan.prefetch, group=group ) as source: eta = group.run(lambda rank: sweep_.run(rank, source))[0] assert eta is not None # noqa: S101 (rank 0 returns the cores) @@ -244,9 +252,7 @@ def get(self, j: int, up: int, down: int) -> np.ndarray: """Return the host tensor ``(chi, up, down)`` of site ``j``.""" with self._lock: if j not in self._drawn: - omega = gaussian_sketch( - self._prng, (self._chi, up, down), self._dtype, np - ) + omega = gaussian_sketch(self._prng, (self._chi, up, down), self._dtype) self._drawn[j] = (omega, [self._ranks]) omega, left = self._drawn[j] left[0] -= 1 @@ -281,19 +287,18 @@ def _tag(self, rank: int) -> str: return "" return f"Device {self.group.devices[rank]}: " - def run(self, rank: int, source: SiteSource) -> list[NDArray] | None: + def run(self, rank: int, source: SiteSource) -> list[np.ndarray] | None: """Run both passes on ``rank``; rank 0 returns the output cores.""" xp, tag = self.group.xp, self._tag(rank) lo, hi = self.columns[rank] env_shapes = [(hi - lo, *(s[0] for s in site)) for site in self.shapes] - kernels = _SplittingKernels(SiteKernels(self.depth), xp, tag) + kernels = _SplittingKernels(SiteKernels(self.depth), tag) with ( device_pool_limit(xp, self.budgets.device_cap), EnvironmentStore( self.plan, env_shapes, self.work, - xp, self.budgets.scratch_dir, copy_stream=new_stream(xp), ) as store, @@ -346,13 +351,13 @@ def _right_to_left( kernels: _SplittingKernels, source: SiteSource, store: EnvironmentStore, - ) -> list[NDArray] | None: + ) -> list[np.ndarray] | None: """Build the output cores, host-side on rank 0, from the last site down.""" xp, plan, group, dtype = self.group.xp, self.plan, self.group, self.work root = rank == 0 col_lo, col_hi = self.columns[rank] n_local = col_hi - col_lo - eta_reversed: list[NDArray] = [] + eta_reversed: list[np.ndarray] = [] proj = xp.ones((1,) * (self.depth + 1), dtype=dtype) for j in range(len(source) - 1, 0, -1): cores = source.next(rank) @@ -372,10 +377,12 @@ def _right_to_left( if root: assert sketch is not None # noqa: S101 (gathered on rank 0) rows = proj.shape[0] * up * down - Q = truncated_qr(sketch.reshape(self.chi_out, rows).T, self.cutoff, xp) + Q = truncated_qr( + xp.reshape(sketch, (self.chi_out, rows)).T, self.cutoff + ) del sketch - eta_full = xp.ascontiguousarray(Q.T).reshape( - -1, proj.shape[0], up, down + eta_full = xp.reshape( + xp.ascontiguousarray(Q.T), (-1, proj.shape[0], up, down) ) del Q eta_local = group.scatter(rank, eta_full) @@ -386,7 +393,9 @@ def _right_to_left( (eta_local.shape[0], *(c.shape[0] for c in cores)), dtype=dtype ) for lo, hi in _ranges(eta_local.shape[0], site.project_batch): - new_local[lo:hi] = kernels.project(eta_local[lo:hi].conj(), cores, proj) + new_local[lo:hi] = kernels.project( + xp.conj(eta_local[lo:hi]), cores, proj + ) del proj, eta_local proj = group.allgather(rank, new_local) del new_local @@ -397,5 +406,7 @@ def _right_to_left( up, down = cores[0].shape[2], cores[-1].shape[3] first = xp.empty((proj.shape[0], up, down), dtype=dtype) for lo, hi in _ranges(proj.shape[0], plan.sites[0].project_batch): - first[lo:hi] = kernels.first(cores, proj[lo:hi]).reshape(hi - lo, up, down) + first[lo:hi] = xp.reshape( + kernels.first(cores, proj[lo:hi]), (hi - lo, up, down) + ) return [to_numpy(first)[None], *reversed(eta_reversed)] diff --git a/src/src_method/_tensor_train.py b/src/src_method/_tensor_train.py index 5e69020..ad8bc95 100644 --- a/src/src_method/_tensor_train.py +++ b/src/src_method/_tensor_train.py @@ -27,7 +27,7 @@ if TYPE_CHECKING: from collections.abc import Sequence - from numpy.typing import NDArray + from .utils import Array # Minimum number of sites for which the randomized SRC sweep is defined. MIN_SRC_SITES = 3 @@ -70,7 +70,7 @@ def __array__(self, dtype: Any = None, copy: bool | None = None) -> np.ndarray: # A site as the public entry points accept it: an array or a lazily read array-like. -Site: TypeAlias = "NDArray | SiteLike" +Site: TypeAlias = "Array | SiteLike" __all__ = [ "MIN_SRC_SITES", @@ -154,8 +154,8 @@ def check_exact_supported(n_sites: int) -> None: def exact_compress( - arrays: Sequence[NDArray], chi_out: int, kind: TrainKind -) -> list[NDArray]: + arrays: Sequence[Array], chi_out: int, kind: TrainKind +) -> list[np.ndarray]: """Compress a two-site train exactly via a single truncated SVD. Site counts are validated by the caller via `check_exact_supported`. @@ -187,7 +187,7 @@ def exact_compress( ] -def _truncated_svd(theta: NDArray, chi_out: int) -> tuple[NDArray, NDArray]: +def _truncated_svd(theta: np.ndarray, chi_out: int) -> tuple[np.ndarray, np.ndarray]: """Split a matrix as ``(U @ diag(S), Vh)``, keeping at most ``chi_out`` values.""" U, S, Vh = np.linalg.svd(theta, full_matrices=False) rank = min(chi_out, S.size) @@ -217,7 +217,7 @@ def padded_shape( return left, right, up, down -def pad_site(site: NDArray, kind: TrainKind, i: int, last: int) -> NDArray: +def pad_site(site: Array, kind: TrainKind, i: int, last: int) -> Array: """View one site as a bulk MPO tensor ``(l, r, u, d)``; see `pad`. Args: @@ -237,7 +237,7 @@ def pad_site(site: NDArray, kind: TrainKind, i: int, last: int) -> NDArray: return view -def pad(train: Sequence[NDArray]) -> list[NDArray]: +def pad(train: Sequence[Array]) -> list[Array]: """View every site of a train as a bulk MPO tensor ``(l, r, u, d)``. Boundary sites gain a size-1 outer bond and MPS sites a size-1 ``d`` leg, so a @@ -255,7 +255,7 @@ def pad(train: Sequence[NDArray]) -> list[NDArray]: return [pad_site(site, kind, i, last) for i, site in enumerate(train)] -def unpad(train: Sequence[NDArray], kind: TrainKind) -> list[NDArray]: +def unpad(train: Sequence[Array], kind: TrainKind) -> list[Array]: """Invert `pad`: drop the size-1 outer bonds and, for an MPS, the ``d`` leg. Args: @@ -398,7 +398,7 @@ def _check_physical_dims( def exact_stack( layers: Sequence[Sequence[Site]], chi_out: int, kind: TrainKind -) -> list[NDArray]: +) -> list[np.ndarray]: """Contract and compress a two-site stack exactly. At each site the layers are folded into one, fusing their bonds, and the @@ -420,7 +420,7 @@ def exact_stack( return exact_compress(unpad(product, kind), chi_out, kind) -def _fuse(upper: NDArray, lower: NDArray) -> NDArray: +def _fuse(upper: np.ndarray, lower: np.ndarray) -> np.ndarray: """Contract ``(l1, r1, u, x) . (l2, r2, x, d)`` into ``(l1 l2, r1 r2, u, d)``.""" l1, r1, up, _ = upper.shape l2, r2, _, down = lower.shape diff --git a/src/src_method/apply.py b/src/src_method/apply.py index 1fd6969..76ff40e 100644 --- a/src/src_method/apply.py +++ b/src/src_method/apply.py @@ -17,7 +17,7 @@ from collections.abc import Sequence import numpy as np - from numpy.typing import DTypeLike, NDArray + from numpy.typing import DTypeLike from ._plan import Resources from ._tensor_train import Site @@ -33,7 +33,7 @@ def apply( seed: int | None = None, device: str = "cpu", resources: Resources | None = None, -) -> list[NDArray]: +) -> list[np.ndarray]: """Applies the Successive Randomized Compression (SRC) algorithm. Equivalent to ``src(left_tensor, right_tensor, ...)`` restricted to an MPO on diff --git a/src/src_method/compress.py b/src/src_method/compress.py index 46e7ffe..a2acb7e 100644 --- a/src/src_method/compress.py +++ b/src/src_method/compress.py @@ -16,7 +16,7 @@ from collections.abc import Sequence import numpy as np - from numpy.typing import DTypeLike, NDArray + from numpy.typing import DTypeLike from ._plan import Resources from ._tensor_train import Site @@ -31,7 +31,7 @@ def compress( seed: int | None = None, device: str = "cpu", resources: Resources | None = None, -) -> list[NDArray]: +) -> list[np.ndarray]: """Applies the Successive Randomized Compression (SRC) algorithm. Equivalent to ``src(tensor, ...)``; see `src_method.stack.src` for the diff --git a/src/src_method/stack.py b/src/src_method/stack.py index 631e314..3452ce7 100644 --- a/src/src_method/stack.py +++ b/src/src_method/stack.py @@ -24,7 +24,7 @@ from collections.abc import Sequence import numpy as np - from numpy.typing import DTypeLike, NDArray + from numpy.typing import DTypeLike from ._plan import Resources from ._tensor_train import Site @@ -45,7 +45,7 @@ def src( seed: int | None = None, device: str = "cpu", resources: Resources | None = None, -) -> list[NDArray]: +) -> list[np.ndarray]: """Contract a stack of tensor trains and compress the result with SRC. The stack ``src(T_1, T_2, ..., T_m)`` is the product ``T_1 T_2 ... T_m`` in @@ -108,7 +108,7 @@ def src( """ chi_out = validate_chi_out(chi_out) validate_cutoff(cutoff) - xp = get_xp(device) + get_xp(device) # Reject an unknown device, or a missing cupy, at once. prng = default_rng(seed) layers, kind = normalize_stack(trains) @@ -123,14 +123,14 @@ def src( n_sites, len(layers), kind, - xp.__name__, + device, ) result = sweep( layers, kind, chi_out, prng, - xp, + device=device, cutoff=cutoff, dtype=sketch_dtype(dtype, *layers), resources=resources, diff --git a/src/src_method/utils/__init__.py b/src/src_method/utils/__init__.py index 54039b8..91b2111 100644 --- a/src/src_method/utils/__init__.py +++ b/src/src_method/utils/__init__.py @@ -3,6 +3,7 @@ from __future__ import annotations from ._backend import ( + Array, NullEvent, NullStream, copy_into, @@ -30,6 +31,7 @@ from .linalg import truncated_qr __all__ = [ + "Array", "NullEvent", "NullStream", "copy_into", diff --git a/src/src_method/utils/_backend.py b/src/src_method/utils/_backend.py index 9a98f79..ff78a37 100644 --- a/src/src_method/utils/_backend.py +++ b/src/src_method/utils/_backend.py @@ -1,11 +1,12 @@ -"""Array-module backend selection for CPU (numpy) and GPU (cupy). - -A single resolver returns the appropriate array module, next to a PRNG factory, -host-transfer helpers and the few stream and memory queries the out-of-core sweep -needs. All hot-loop code paths receive an ``xp`` module and call ``xp.linalg.*`` / -``xp.asarray`` directly, so backend selection adds zero per-op overhead. On the -host backend the stream helpers are synchronous no-ops, so the same algorithm code -runs on both. +"""The seam between the algorithms and the array backends: NumPy and CuPy. + +The algorithms speak the array API standard through `array_api_compat`, and take +the namespace from their arrays (``array_namespace(x)``). Arrays stay plain ``numpy.ndarray`` or ``cupy.ndarray``, +so `opt_einsum` dispatches on them as before. What the standard leaves out lives +here: picking the namespace for a ``device``, streams and events, page-locked +memory, CuPy's memory pool, and the host-side random draws. Functions that create +arrays from nothing take the namespace ``xp``; on the host the stream helpers are +synchronous no-ops, so the same algorithm code runs on both backends. """ from __future__ import annotations @@ -13,38 +14,43 @@ import os from contextlib import contextmanager from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, TypeAlias import numpy as np +from array_api_compat import numpy as np_compat if TYPE_CHECKING: from collections.abc import Iterator, Sequence from types import ModuleType - from numpy.typing import DTypeLike, NDArray + from numpy.typing import DTypeLike from src_method._tensor_train import Site, SiteLike +# An array of the active backend: a numpy.ndarray or a cupy.ndarray. CuPy ships no +# type stubs, so to a type checker either is no narrower than Any. +Array: TypeAlias = Any + def get_xp(device: str) -> ModuleType: - """Return the array module for the requested device. + """Return the array API namespace for the requested device. Args: device: ``"cpu"`` for numpy or ``"gpu"`` for cupy. Returns: - The numpy or cupy module. + The `array_api_compat` namespace of numpy or cupy. Raises: ValueError: If ``device`` is not recognised. ImportError: If ``device="gpu"`` but cupy is not installed. """ if device == "cpu": - return np + return np_compat if device == "gpu": - import cupy # noqa: PLC0415 (lazy: optional dependency) + import array_api_compat.cupy as xp # noqa: PLC0415 (optional dependency) - return cupy + return xp msg = f"Unknown device {device!r}; expected 'cpu' or 'gpu'." raise ValueError(msg) @@ -61,7 +67,7 @@ def default_rng( return np.random.default_rng(seed) -def to_numpy(arr: NDArray | SiteLike) -> np.ndarray: +def to_numpy(arr: Array | SiteLike) -> np.ndarray: """Bring an array, or a lazily read site, onto the host as a numpy array. A no-op for numpy arrays. @@ -96,7 +102,8 @@ def synchronize(self) -> None: def is_host(xp: ModuleType) -> bool: """Whether ``xp`` is the host backend (numpy).""" - return xp is np + # An identity check: the stream helpers call this several times per site. + return xp is np_compat or xp is np def new_stream(xp: ModuleType) -> Any: # noqa: ANN401 (a cupy or null stream) @@ -118,7 +125,7 @@ def pinned_empty(n_bytes: int, xp: ModuleType) -> np.ndarray: Args: n_bytes: The size of the buffer. - xp: Array module (``numpy`` or ``cupy``). + xp: The array namespace. Returns: A ``uint8`` host array of ``n_bytes`` elements. @@ -130,7 +137,7 @@ def pinned_empty(n_bytes: int, xp: ModuleType) -> np.ndarray: return cupyx.empty_pinned(n_bytes, dtype=np.uint8) -def to_device_async(host: np.ndarray, xp: ModuleType, stream: Any) -> NDArray: # noqa: ANN401 +def to_device_async(host: np.ndarray, xp: ModuleType, stream: Any) -> Array: # noqa: ANN401 """Copy a host array to a new device array on ``stream``. On the host backend this is a plain copy, so the result never aliases a staging @@ -138,7 +145,7 @@ def to_device_async(host: np.ndarray, xp: ModuleType, stream: Any) -> NDArray: Args: host: The source; page-locked for the copy to be asynchronous. - xp: Array module (``numpy`` or ``cupy``). + xp: The array namespace of the device. stream: The stream that performs the copy. Returns: @@ -151,7 +158,7 @@ def to_device_async(host: np.ndarray, xp: ModuleType, stream: Any) -> NDArray: return device -def to_host_async(device: NDArray, out: np.ndarray, stream: Any) -> None: # noqa: ANN401 +def to_host_async(device: Array, out: np.ndarray, stream: Any) -> None: # noqa: ANN401 """Copy a device array into a host array on ``stream``. The copy is complete once an event recorded on ``stream`` afterwards is. @@ -193,7 +200,7 @@ def use_device(xp: ModuleType, device: int) -> Iterator[None]: so this is a no-op. Args: - xp: Array module (``numpy`` or ``cupy``). + xp: The array namespace. device: The device id. Yields: @@ -216,7 +223,7 @@ def enable_peer_access(xp: ModuleType, devices: Sequence[int]) -> bool: Peer copies work without it, but are then staged through the host. Args: - xp: Array module (``numpy`` or ``cupy``). + xp: The array namespace. devices: The device ids. Returns: @@ -242,7 +249,7 @@ def enable_peer_access(xp: ModuleType, devices: Sequence[int]) -> bool: return complete -def copy_into(dst: NDArray, src: NDArray, xp: ModuleType, stream: Any) -> None: # noqa: ANN401 +def copy_into(dst: Array, src: Array, stream: Any) -> None: # noqa: ANN401 """Copy ``src`` into ``dst``, possibly on another device, on ``stream``. Both arrays must be C-contiguous with the same shape and dtype. Across devices @@ -254,7 +261,6 @@ def copy_into(dst: NDArray, src: NDArray, xp: ModuleType, stream: Any) -> None: Args: dst: The destination. src: The source. - xp: Array module (``numpy`` or ``cupy``). stream: The stream that performs the copy. """ if dst.shape != src.shape or dst.dtype != src.dtype: @@ -262,31 +268,31 @@ def copy_into(dst: NDArray, src: NDArray, xp: ModuleType, stream: Any) -> None: raise ValueError(msg) if dst.nbytes == 0: return - if is_host(xp): + if isinstance(dst, np.ndarray): np.copyto(dst, src) return if not (dst.flags.c_contiguous and src.flags.c_contiguous): msg = "Device copies need C-contiguous arrays." raise ValueError(msg) + import cupy # noqa: PLC0415 (optional dependency) + # CuPy arrays: ``data`` is a memory pointer and ``device`` a CUDA device. - to: Any = dst - source: Any = src - runtime = xp.cuda.runtime - if to.device.id == source.device.id: + runtime = cupy.cuda.runtime + if dst.device.id == src.device.id: runtime.memcpyAsync( - to.data.ptr, - source.data.ptr, - to.nbytes, + dst.data.ptr, + src.data.ptr, + dst.nbytes, runtime.memcpyDeviceToDevice, stream.ptr, ) else: runtime.memcpyPeerAsync( - to.data.ptr, - to.device.id, - source.data.ptr, - source.device.id, - to.nbytes, + dst.data.ptr, + dst.device.id, + src.data.ptr, + src.device.id, + dst.nbytes, stream.ptr, ) @@ -298,7 +304,7 @@ def device_memory(xp: ModuleType) -> tuple[int, int]: by cupy's default memory pool. Args: - xp: The cupy module. + xp: The array namespace of cupy. Returns: ``(available, total)`` in bytes. @@ -337,7 +343,7 @@ def device_pool_limit(xp: ModuleType, budget: int | None) -> Iterator[None]: without a budget. Args: - xp: Array module (``numpy`` or ``cupy``). + xp: The array namespace. budget: The bytes the pool may allocate on top of those in use, or ``None`` for no cap. @@ -396,32 +402,28 @@ def sketch_dtype(dtype: DTypeLike | None, *inputs: Sequence[Site]) -> np.dtype: def gaussian_sketch( - prng: np.random.Generator, - shape: tuple[int, ...], - dtype: DTypeLike, - xp: ModuleType, -) -> NDArray: - """Draw a Gaussian test tensor of the given dtype on the device of ``xp``. + prng: np.random.Generator, shape: tuple[int, ...], dtype: DTypeLike +) -> np.ndarray: + """Draw a Gaussian test tensor of the given dtype on the host. Real dtypes get i.i.d. standard normal entries. Complex dtypes get the complex Ginibre ensemble, ``(x + iy) / sqrt(2)`` with ``x, y`` i.i.d. standard normal, whose law is invariant under unitary transformations; a real sketch of a complex operator is not, and loses the guarantees of randomized range finding. - The draw is always made on the host in double precision and then cast, so a - seed gives the same sketch on every device and at every precision. + The draw is made in double precision and then cast, so a seed gives the same + sketch on every device and at every precision. Args: prng: Host-side random number generator. shape: Shape of the tensor. dtype: Dtype of the tensor. - xp: Array module (``numpy`` or ``cupy``). Returns: - The sketch as an array of ``xp``. + The sketch. """ dtype = np.dtype(dtype) if dtype.kind == "c": re, im = prng.normal(size=(2, *shape)) / np.sqrt(2.0) - return xp.asarray(re + 1j * im).astype(dtype) - return xp.asarray(prng.normal(size=shape)).astype(dtype) + return (re + 1j * im).astype(dtype) + return prng.normal(size=shape).astype(dtype) diff --git a/src/src_method/utils/linalg.py b/src/src_method/utils/linalg.py index ff2e89c..b8d00f0 100644 --- a/src/src_method/utils/linalg.py +++ b/src/src_method/utils/linalg.py @@ -5,18 +5,15 @@ from typing import TYPE_CHECKING import numpy as np +from array_api_compat import array_namespace -if TYPE_CHECKING: - from types import ModuleType +from ._backend import to_numpy - from numpy.typing import NDArray +if TYPE_CHECKING: + from ._backend import Array -def truncated_qr( - matrix: NDArray[np.floating | np.complexfloating], - cutoff: float, - xp: ModuleType = np, -) -> NDArray[np.floating | np.complexfloating]: +def truncated_qr(matrix: Array, cutoff: float) -> Array: """QR with SVD-based rank truncation, returning only the isometry. Decomposes ``matrix = Q @ R``, then truncates via SVD on R, @@ -25,31 +22,23 @@ def truncated_qr( When ``cutoff <= 0`` this falls back to a plain QR (no truncation). Args: - matrix: Input matrix of shape (m, n). + matrix: Input matrix of shape (m, n), on any backend. cutoff: Relative singular-value threshold. Singular values satisfying ``s < cutoff * s_max`` are discarded. - xp: Array module (``numpy`` or ``cupy``); defaults to numpy. Returns: - The truncated isometry Q of shape ``(m, rank)``. + The truncated isometry Q of shape ``(m, rank)``, on the device of + ``matrix``. """ - m, n = matrix.shape - transpose = m < n - + xp = array_namespace(matrix) if cutoff <= 0: return xp.linalg.qr(matrix)[0] + transpose = matrix.shape[0] < matrix.shape[1] Q, R = xp.linalg.qr(matrix.T if transpose else matrix) - R_np = R.get() if hasattr(R, "get") else np.asarray(R) - U, S, _ = np.linalg.svd(R_np.T if transpose else R_np, full_matrices=False) + # R is small: its SVD runs on the host. + R = to_numpy(R) + U, S, _ = np.linalg.svd(R.T if transpose else R, full_matrices=False) rank = max(1, int((cutoff * S[0] <= S).sum())) - - if transpose: - Q_trunc = U[:, :rank] - else: - U_trunc = xp.asarray(U[:, :rank]) if xp is not np else U[:, :rank] - Q_trunc = Q @ U_trunc - - if xp is not np: - return xp.asarray(Q_trunc) - return Q_trunc + U = xp.asarray(U[:, :rank]) + return U if transpose else Q @ U diff --git a/tests/test_backend.py b/tests/test_backend.py index d467961..62a1ca1 100644 --- a/tests/test_backend.py +++ b/tests/test_backend.py @@ -13,6 +13,7 @@ device_pool_bytes, device_pool_limit, enable_peer_access, + get_xp, host_memory_available, is_host, new_stream, @@ -27,6 +28,7 @@ def test_host_streams_are_null(): stream = new_stream(np) assert is_host(np) + assert is_host(get_xp("cpu")) assert isinstance(stream, NullStream) assert isinstance(current_stream(np), NullStream) event = stream.record() @@ -91,13 +93,13 @@ def test_host_devices_are_simulated(): def test_copy_into_on_host(): dst = np.empty((2, 3), dtype=complex) - copy_into(dst, np.arange(6.0).reshape(2, 3).astype(complex), np, NullStream()) + copy_into(dst, np.arange(6.0).reshape(2, 3).astype(complex), NullStream()) np.testing.assert_array_equal(dst, np.arange(6.0).reshape(2, 3)) def test_copy_into_rejects_mismatches(): with pytest.raises(ValueError, match="Cannot copy"): - copy_into(np.empty(3), np.empty(4), np, NullStream()) + copy_into(np.empty(3), np.empty(4), NullStream()) with pytest.raises(ValueError, match="Cannot copy"): - copy_into(np.empty(3), np.empty(3, dtype=np.float32), np, NullStream()) + copy_into(np.empty(3), np.empty(3, dtype=np.float32), NullStream()) diff --git a/tests/test_group.py b/tests/test_group.py index c0bc6a0..bb3dc0a 100644 --- a/tests/test_group.py +++ b/tests/test_group.py @@ -21,7 +21,7 @@ def test_blocks_partition_in_rank_order(n, size): def test_a_group_of_one_runs_inline(): - group = DeviceGroup(np, [0]) + group = DeviceGroup("cpu", [0]) local = np.arange(6.0).reshape(3, 2) def fn(rank): @@ -37,7 +37,7 @@ def fn(rank): @pytest.mark.parametrize("size", [2, 3]) @pytest.mark.parametrize("dtype", [np.float64, np.complex128]) def test_collectives(size, dtype): - group = DeviceGroup(np, list(range(size))) + group = DeviceGroup("cpu", size) # Uneven blocks: rank r holds r + 1 rows. parts = [ (np.arange((r + 1) * 6) + 100 * r).reshape(r + 1, 2, 3) for r in range(size) @@ -65,7 +65,7 @@ def fn(rank): def test_many_collectives_in_a_row(): - group = DeviceGroup(np, [0, 1, 2]) + group = DeviceGroup("cpu", [0, 1, 2]) def fn(rank): total = np.zeros(3) @@ -78,7 +78,7 @@ def fn(rank): def test_a_failing_rank_releases_the_others(): - group = DeviceGroup(np, [0, 1, 2]) + group = DeviceGroup("cpu", [0, 1, 2]) def fn(rank): if rank == 1: @@ -93,7 +93,7 @@ def fn(rank): def test_abort_wakes_registered_waiters(): - group = DeviceGroup(np, [0, 1]) + group = DeviceGroup("cpu", [0, 1]) released = threading.Event() group.on_abort(released.set) diff --git a/tests/test_package.py b/tests/test_package.py index 3a9e326..410ede4 100644 --- a/tests/test_package.py +++ b/tests/test_package.py @@ -753,7 +753,7 @@ def test_stack_dtype_follows_inputs() -> None: @pytest.mark.parametrize("dtype", [np.float32, np.float64, np.complex64, np.complex128]) def test_gaussian_sketch(dtype) -> None: prng = np.random.default_rng(0) - omega = gaussian_sketch(prng, (400, 400), dtype, np) + omega = gaussian_sketch(prng, (400, 400), dtype) assert omega.dtype == dtype assert omega.var() == pytest.approx(1.0, rel=0.05) @@ -765,6 +765,6 @@ def test_gaussian_sketch(dtype) -> None: def test_gaussian_sketch_seed_is_precision_independent() -> None: - single = gaussian_sketch(np.random.default_rng(3), (5, 5), np.complex64, np) - double = gaussian_sketch(np.random.default_rng(3), (5, 5), np.complex128, np) + single = gaussian_sketch(np.random.default_rng(3), (5, 5), np.complex64) + double = gaussian_sketch(np.random.default_rng(3), (5, 5), np.complex128) np.testing.assert_allclose(single, double, rtol=1e-6) diff --git a/tests/test_plan.py b/tests/test_plan.py index 67fb38c..6c4b86b 100644 --- a/tests/test_plan.py +++ b/tests/test_plan.py @@ -2,13 +2,13 @@ from contextlib import contextmanager from pathlib import Path -from types import ModuleType import numpy as np import pytest import src_method._plan as plan_module from src_method import Resources +from src_method._group import DeviceGroup from src_method._plan import ( GEMM_MULTIPLE, Budgets, @@ -18,6 +18,7 @@ ) GB = 10**9 +HOST = DeviceGroup("cpu") @pytest.mark.parametrize( @@ -53,7 +54,7 @@ def test_resources_validate_at_construction(): def test_resolve_budgets_explicit_on_host(tmp_path): - budgets = resolve_budgets(Resources(host_memory="2GB", scratch_dir=tmp_path), np) + budgets = resolve_budgets(Resources(host_memory="2GB", scratch_dir=tmp_path), HOST) assert budgets.unified assert budgets.host == budgets.device == 2 * GB @@ -64,7 +65,7 @@ def test_resolve_budgets_explicit_on_host(tmp_path): def test_resolve_budgets_detects_host_memory(monkeypatch, tmp_path): monkeypatch.setattr(plan_module, "host_memory_available", lambda: 10 * GB) - budgets = resolve_budgets(Resources(scratch_dir=tmp_path / "not" / "yet"), np) + budgets = resolve_budgets(Resources(scratch_dir=tmp_path / "not" / "yet"), HOST) assert budgets.host == 9 * GB @@ -91,11 +92,7 @@ def use_device(_xp, device): @pytest.mark.usefixtures("fake_gpus") def test_resolve_budgets_detects_device_memory(tmp_path): - fake_xp = ModuleType("fake_xp") - - budgets = resolve_budgets( - Resources(host_memory="1GB", scratch_dir=tmp_path), fake_xp - ) + budgets = resolve_budgets(Resources(host_memory="1GB", scratch_dir=tmp_path), HOST) assert not budgets.unified assert budgets.device == 26 * GB # minus max(10% of 40 GB, 1 GiB) @@ -104,11 +101,8 @@ def test_resolve_budgets_detects_device_memory(tmp_path): @pytest.mark.usefixtures("fake_gpus") def test_resolve_budgets_caps_explicit_device_memory_with_margin(tmp_path): - fake_xp = ModuleType("fake_xp") - budgets = resolve_budgets( - Resources(gpu_memory="10GB", host_memory="1GB", scratch_dir=tmp_path), - fake_xp, + Resources(gpu_memory="10GB", host_memory="1GB", scratch_dir=tmp_path), HOST ) assert budgets.device == 10 * GB @@ -119,9 +113,7 @@ def test_resolve_budgets_takes_the_smallest_device(fake_gpus, tmp_path): fake_gpus[1] = (20 * GB, 40 * GB) budgets = resolve_budgets( - Resources(host_memory="1GB", scratch_dir=tmp_path), - ModuleType("fake_xp"), - devices=(0, 1), + Resources(host_memory="1GB", scratch_dir=tmp_path), DeviceGroup("cpu", 2) ) assert budgets.device == 16 * GB @@ -130,7 +122,7 @@ def test_resolve_budgets_takes_the_smallest_device(fake_gpus, tmp_path): def test_resolve_budgets_splits_the_host_among_simulated_devices(tmp_path): budgets = resolve_budgets( - Resources(host_memory="8GB", scratch_dir=tmp_path), np, devices=(0, 1, 2, 3) + Resources(host_memory="8GB", scratch_dir=tmp_path), DeviceGroup("cpu", 4) ) assert budgets.host == 8 * GB @@ -138,7 +130,7 @@ def test_resolve_budgets_splits_the_host_among_simulated_devices(tmp_path): def test_resolve_budgets_leaves_the_host_uncapped(tmp_path): - budgets = resolve_budgets(Resources(host_memory="1GB", scratch_dir=tmp_path), np) + budgets = resolve_budgets(Resources(host_memory="1GB", scratch_dir=tmp_path), HOST) assert budgets.device_cap is None @@ -147,7 +139,7 @@ def test_resolve_budgets_detects_disk(monkeypatch, tmp_path): usage = type("Usage", (), {"free": 100 * GB}) monkeypatch.setattr(plan_module.shutil, "disk_usage", lambda _path: usage) - budgets = resolve_budgets(Resources(host_memory="1GB", scratch_dir=tmp_path), np) + budgets = resolve_budgets(Resources(host_memory="1GB", scratch_dir=tmp_path), HOST) assert budgets.disk == 95 * GB diff --git a/tests/test_sites.py b/tests/test_sites.py index 1f45669..07a65e1 100644 --- a/tests/test_sites.py +++ b/tests/test_sites.py @@ -67,7 +67,9 @@ def test_sites_come_in_the_given_order(depth): rng = np.random.default_rng(1) train = CountingTrain(random_mpo(4, 3, rng)) - with SiteSource([train], np, [2, 3, 2], depth=depth) as source: + with SiteSource( + [train], [2, 3, 2], depth=depth, group=DeviceGroup("cpu") + ) as source: cores = [source.next() for _ in range(3)] assert train.reads == [0, 0, 2, 1] @@ -80,10 +82,10 @@ def test_simulated_devices_share_one_read(depth): rng = np.random.default_rng(1) train = CountingTrain(random_mpo(4, 3, rng)) mps = [t[..., 0] for t in random_mpo(4, 2, rng)] - group = DeviceGroup(np, [0, 1, 2]) + group = DeviceGroup("cpu", [0, 1, 2]) order = sweep_order(4) - with SiteSource([train, mps], np, order, depth=depth, group=group) as source: + with SiteSource([train, mps], order, depth=depth, group=group) as source: got = group.run(lambda rank: [source.next(rank) for _ in order]) assert train.reads == [2, 2, 2, 1] @@ -104,12 +106,12 @@ def __getitem__(self, j): return Broken(self, j) if j == 1 else CountingSite(self, j) rng = np.random.default_rng(1) - group = DeviceGroup(np, [0, 1]) + group = DeviceGroup("cpu", [0, 1]) order = sweep_order(4) with ( SiteSource( - [BrokenTrain(random_mpo(4, 3, rng))], np, order, depth=1, group=group + [BrokenTrain(random_mpo(4, 3, rng))], order, depth=1, group=group ) as source, pytest.raises(OSError, match="unreadable"), ): diff --git a/tests/test_stack.py b/tests/test_stack.py index 404caa5..4bce1c1 100644 --- a/tests/test_stack.py +++ b/tests/test_stack.py @@ -103,7 +103,7 @@ def test_sweep_is_exact_without_truncation(spec, rng): stack = make_stack(spec, rng) kind = "mps" if stack[-1][0].ndim == 2 else "mpo" - out = sweep(stack, kind, CHI_EXACT, np.random.default_rng(0), np, dtype=np.float64) + out = sweep(stack, kind, CHI_EXACT, np.random.default_rng(0), dtype=np.float64) assert rel_error(out, dense_stack(*stack)) < 1e-10 @@ -114,7 +114,7 @@ def test_sweep_rectangular_physical_legs(rng): psi = random_mps([2, 2, 3, 2], rng, phys=2) out = sweep( - [A, B, psi], "mps", CHI_EXACT, np.random.default_rng(0), np, dtype=np.float64 + [A, B, psi], "mps", CHI_EXACT, np.random.default_rng(0), dtype=np.float64 ) assert rel_error(out, dense_stack(A, B, psi)) < 1e-10 @@ -124,7 +124,7 @@ def test_sweep_does_not_mutate_inputs(rng): stack = make_stack("A B psi", rng) before = [[t.copy() for t in train] for train in stack] - sweep(stack, "mps", 4, np.random.default_rng(0), np, dtype=np.float64) + sweep(stack, "mps", 4, np.random.default_rng(0), dtype=np.float64) for train, saved in zip(stack, before): for t, s in zip(train, saved): diff --git a/tests/test_store.py b/tests/test_store.py index c9a22a7..6826519 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -41,7 +41,6 @@ def open_store(tier, tmp_path, **batches): make_plan(tier, **batches), SHAPES, np.complex128, - np, tmp_path, copy_stream=NullStream(), ) diff --git a/uv.lock b/uv.lock index b68e420..26f8825 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 3 +revision = 5 requires-python = ">=3.11, <3.15" resolution-markers = [ "python_full_version >= '3.14'", @@ -25,6 +25,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/81/29/5ecc3a15d5a33e31b26c11426c45c501e439cb865d0bff96315d86443b78/appnope-0.1.4-py2.py3-none-any.whl", hash = "sha256:502575ee11cd7a28c0205f379b525beefebab9d161b7c964670864014ed7213c", size = 4321, upload-time = "2024-02-06T09:43:09.663Z" }, ] +[[package]] +name = "array-api-compat" +version = "1.15.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/86/12/65d297d3e425cb7fe5247b7ec14b4ed83539e9b5e89b65c15ea7ee274182/array_api_compat-1.15.0.tar.gz", hash = "sha256:53c5f922491bf15f62847afafc4e39eedfae57d218988fefb8cce39c2a9b3dea", size = 129305, upload-time = "2026-06-07T20:53:24.964Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/86/16/1a8fd2b19544b84575cf84ef7aa3ad4c173b756d5f087c91f85d1b295777/array_api_compat-1.15.0-py3-none-any.whl", hash = "sha256:7b1b9c53269061403fd5f45a8de349f16e7887653328bfa0c5f2d45299ff0a8e", size = 79113, upload-time = "2026-06-07T20:53:23.621Z" }, +] + [[package]] name = "asttokens" version = "3.0.1" @@ -2433,6 +2442,7 @@ wheels = [ name = "src-method" source = { editable = "." } dependencies = [ + { name = "array-api-compat" }, { name = "numpy" }, { name = "opt-einsum" }, ] @@ -2497,6 +2507,7 @@ test = [ [package.metadata] requires-dist = [ + { name = "array-api-compat", specifier = ">=1.12" }, { name = "cupy", marker = "extra == 'gpu-rocm'", specifier = ">=14.1.1" }, { name = "cupy-cuda13x", extras = ["ctk"], marker = "extra == 'gpu-nvidia'", specifier = ">=14" }, { name = "numpy", specifier = ">=2.0,<3" },