diff --git a/README.md b/README.md index 70292b3..fa97463 100644 --- a/README.md +++ b/README.md @@ -87,6 +87,26 @@ The result then pairs with a ket by plain contraction, with no further conjugati At runtime, pass `device="gpu"` to use GPU acceleration. The library handles backend dispatch automatically. +### Large Problems + +Stacks that do not fit on the GPU, or in host memory, still run: the sweep reads the input cores one site at a time (from NumPy arrays, `np.memmap`, zarr or HDF5 datasets), batches every contraction to a memory budget, and keeps the sketched environments on the GPU, in host memory or on local disk. The budgets are detected, or set explicitly with `Resources`: + +```python +from src_method import Resources, src + +out = src( + N, + V, + M, + U, + chi_out=2000, + device="gpu", + resources=Resources(gpu_memory="36GB", scratch_dir="/local/scratch"), +) +``` + +See [Large problems](https://algorithmiq.github.io/src-method/features/large-problems) for the budgets, the scratch directory and how to read the logged plan. + ### Logging `src_method` logs through the standard library `logging` module, under the @@ -108,7 +128,7 @@ logging.getLogger("src_method").setLevel(logging.DEBUG) # CPU only (default) uv pip install src_method -# With NVIDIA GPU support (CUDA 12.x) +# With NVIDIA GPU support (CUDA 13.x, driver >= 580) uv pip install "src_method[gpu-nvidia]" # With AMD GPU support (ROCm) diff --git a/benches/README.md b/benches/README.md index 61b79ce..8a49f83 100644 --- a/benches/README.md +++ b/benches/README.md @@ -10,6 +10,9 @@ different SRC variants. They are grouped by the kind of workload exercised: - [`stack/`](stack/) — One-shot SRC over stacks of trains against sequential pairwise application, by stack depth: accuracy against dense references and wall time. See [`stack/README.md`](stack/README.md). +- [`large/`](large/) — Out-of-core SRC of `N . V . M . U` with a large `M` read + from disk, on one GPU: plan, wall time, memory peaks and stall time. See + [`large/README.md`](large/README.md). The scripts need the `bench` dependency group (included in `dev`): diff --git a/benches/large/README.md b/benches/large/README.md new file mode 100644 index 0000000..24e15f4 --- /dev/null +++ b/benches/large/README.md @@ -0,0 +1,33 @@ +# Out-of-core benchmark + +`bench_large.py` compresses the stack `N . V . M . U` of Pauli-transfer-matrix MPOs +(physical legs of 4) on one GPU, where `M` has a large bond and is read from disk. + +```bash +# Write the stack to node-local NVMe (205 GB for M at the defaults, complex128). +uv run python benches/large/bench_large.py generate /local/stack --bond-m 4000 + +# Compress it with detected budgets, spilling to the same disk. +uv run python benches/large/bench_large.py run /local/stack --chi-out 2000 \ + --scratch-dir /local/scratch + +# Check that the plan does not change the result: two GPU budgets, same seed. +uv run python benches/large/bench_large.py generate /local/medium --bond-m 1000 +uv run python benches/large/bench_large.py compare /local/medium --chi-out 500 \ + --small 40GB --large 80GB --scratch-dir /local/scratch +``` + +`run` logs the wall time, the size of CuPy's pool at the end (its high-water mark, +since the pool keeps its blocks), the host peak (`ru_maxrss`), the planned peaks, +the bytes spilled and the tiers. With `--debug` before the command, `src_method` also +logs its plan, the time of each pass and the stall time (`SRC stalls`). + +Acceptance of the out-of-core sweep on one A100-40GB, with at least 300 GB of host +memory and node-local NVMe: + +1. `run` completes at 50 sites, `D_M = 4000`, `chi_out = 2000`, complex128. +2. The pool size stays within the GPU budget plus the `max(10%, 1 GiB)` margin + the pool cap allows for fragmentation, and the host peak within the host + budget. +3. The stall time is below 10% of the wall time. +4. `compare` at `D_M = 1000`, `chi_out = 500` reports a distance below `1e-10`. diff --git a/benches/large/bench_large.py b/benches/large/bench_large.py new file mode 100644 index 0000000..8ad4143 --- /dev/null +++ b/benches/large/bench_large.py @@ -0,0 +1,215 @@ +"""Out-of-core SRC of ``N . V . M . U`` with a large ``M``, read from disk. + +Three commands: + +- ``generate``: write the four MPOs site by site as ``.npy`` files, so that ``M`` + is never whole in memory. +- ``run``: compress the stack with automatic (or given) budgets and report the + plan, the wall time, the device pool size, the host peak and the stall time. +- ``compare``: run twice with two GPU budgets and report the relative distance + between the two outputs, which checks that the plan does not change the result. + +The plan, the per-pass times and the stall time are also logged by `src_method` +itself at ``DEBUG``; pass ``--debug`` to see them. +""" + +from __future__ import annotations + +import logging +import resource +from pathlib import Path # noqa: TC003 (cyclopts reads the annotations at runtime) +from time import perf_counter +from typing import Annotated + +import cyclopts +import numpy as np + +import src_method._sweep as sweep_module +from src_method import Resources, src +from src_method._plan import make_plan + +logger = logging.getLogger(__name__) +app = cyclopts.App(help="Benchmark out-of-core SRC of N.V.M.U with a large M.") + +LAYERS = ("N", "V", "M", "U") +PHYS = 4 # Pauli transfer matrix legs + + +def _site_shape(j: int, n_sites: int, bond: int) -> tuple[int, ...]: + left = () if j == 0 else (bond,) + right = () if j == n_sites - 1 else (bond,) + return (*left, *right, PHYS, PHYS) + + +@app.command +def generate( + directory: Path, + *, + n_sites: int = 50, + bond: int = 4, + bond_m: int = 4000, + dtype: str = "complex128", + seed: int = 0, +) -> None: + """Write random MPOs ``N``, ``V``, ``M``, ``U`` as one ``.npy`` file per site. + + Args: + directory: Where to write, ideally node-local NVMe. + n_sites: Number of sites. + bond: Bond dimension of ``N``, ``V`` and ``U``. + bond_m: Bond dimension of ``M``. + dtype: ``float64`` or ``complex128``. + seed: Seed of the random draws. + """ + rng = np.random.default_rng(seed) + kind = np.dtype(dtype) + for name in LAYERS: + chi = bond_m if name == "M" else bond + (directory / name).mkdir(parents=True, exist_ok=True) + # Scaled so that products of the layers stay of order one. + scale = 1 / np.sqrt(chi * PHYS) + for j in range(n_sites): + shape = _site_shape(j, n_sites, chi) + site = np.lib.format.open_memmap( + directory / name / f"{j:04d}.npy", mode="w+", dtype=kind, shape=shape + ) + for row in range(shape[0]): # one left-bond slice at a time + draw = rng.normal(size=shape[1:]) + if kind.kind == "c": + draw = draw + 1j * rng.normal(size=shape[1:]) + site[row] = draw * scale + site.flush() + del site + logger.info("Layer %s written: bond %d", name, chi) + + +def _load(directory: Path) -> list[list[np.ndarray]]: + return [ + [ + np.load(path, mmap_mode="r") + for path in sorted((directory / name).glob("*.npy")) + ] + for name in LAYERS + ] + + +def _run( + directory: Path, chi_out: int, resources: Resources, seed: int +) -> list[np.ndarray]: + layers = _load(directory) + plans = [] + + def spy(*args: object, **kwargs: object) -> object: + plans.append(make_plan(*args, **kwargs)) + return plans[-1] + + sweep_module.make_plan = spy # record the plan of the run + try: + start = perf_counter() + out = src( + *layers, + chi_out=chi_out, + dtype=layers[0][0].dtype, + seed=seed, + device="gpu", + resources=resources, + ) + seconds = perf_counter() - start + finally: + sweep_module.make_plan = make_plan + import cupy # noqa: PLC0415 (GPU-only benchmark) + + (plan,) = plans + tiers = [site.tier for site in plan.sites] + logger.info( + "Run complete: %.1f s, pool %d B, host peak %d B, planned device peak %d B, " + "planned host peak %d B, disk %d B, tiers %s, sketch batches %s", + seconds, + cupy.get_default_memory_pool().total_bytes(), + resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * 1024, + plan.device_peak, + plan.host_peak, + plan.disk_bytes, + {tier: tiers.count(tier) for tier in ("device", "host", "disk")}, + sorted({site.sketch_batch for site in plan.sites[1:]}), + ) + return out + + +@app.command +def run( + directory: Path, + *, + chi_out: int = 2000, + gpu_memory: str | None = None, + host_memory: str | None = None, + scratch_dir: Path | None = None, + seed: int = 0, +) -> None: + """Compress the stack written by ``generate`` on the GPU. + + Args: + directory: The directory given to ``generate``. + chi_out: The output bond dimension. + gpu_memory: GPU budget, e.g. ``36GB``; detected when omitted. + host_memory: Host budget; detected when omitted. + scratch_dir: Where to spill environments; ``$TMPDIR`` when omitted. + seed: Seed of the sketch. + """ + _run(directory, chi_out, Resources(gpu_memory, host_memory, scratch_dir), seed) + + +def _inner(a: list[np.ndarray], b: list[np.ndarray]) -> complex: + """Frobenius inner product ```` of two MPOs, site by site.""" + env = np.einsum("rud,sud->rs", a[0].conj(), b[0]) + for x, y in zip(a[1:-1], b[1:-1]): + env = np.einsum("rs,rtud,svud->tv", env, x.conj(), y, optimize=True) + return complex(np.einsum("rs,rud,sud->", env, a[-1].conj(), b[-1])) + + +@app.command +def compare( + directory: Path, + *, + chi_out: int = 500, + small: str = "40GB", + large: str = "80GB", + scratch_dir: Path | None = None, + seed: int = 0, +) -> None: + """Run with two GPU budgets and report the relative distance of the outputs. + + Args: + directory: The directory given to ``generate``. + chi_out: The output bond dimension. + small: The smaller GPU budget. + large: The larger GPU budget. + scratch_dir: Where to spill environments; ``$TMPDIR`` when omitted. + seed: Seed of the sketch, the same for both runs. + """ + first = _run(directory, chi_out, Resources(small, None, scratch_dir), seed) + second = _run(directory, chi_out, Resources(large, None, scratch_dir), seed) + aa, bb, ab = _inner(first, first), _inner(second, second), _inner(first, second) + distance = np.sqrt(max((aa + bb - 2 * ab).real, 0.0) / aa.real) + logger.info("Relative distance between the runs: %.3e", distance) + + +@app.meta.default +def main( + *tokens: Annotated[str, cyclopts.Parameter(show=False, allow_leading_hyphen=True)], + debug: bool = False, +) -> None: + """Configure logging, then dispatch to a command. + + Args: + tokens: The command and its arguments. + debug: Also show the `src_method` debug log: plan, pass times and stalls. + """ + logging.basicConfig(level=logging.INFO, format="%(message)s") + if debug: + logging.getLogger("src_method").setLevel(logging.DEBUG) + app(tokens) + + +if __name__ == "__main__": + app.meta() diff --git a/docs/content/docs/contributing/testing.mdx b/docs/content/docs/contributing/testing.mdx index d950202..02e5b48 100644 --- a/docs/content/docs/contributing/testing.mdx +++ b/docs/content/docs/contributing/testing.mdx @@ -18,6 +18,11 @@ Always run tests through `pytest`, never as `python tests/test_.py`. | `test_package.py` | `apply` and `compress` against `quimb` references, validation, precision | | `test_stack.py` | `src` over stacks and the sweep behind it | | `test_tensor_train.py` | the shared train helpers in `_tensor_train` | +| `test_backend.py` | the stream, staging and memory helpers in `utils._backend` | +| `test_kernels.py` | the batched site contractions and their peak-memory estimate | +| `test_plan.py` | the planner: budgets, batch sizes and environment tiers | +| `test_store.py` | environment storage on the device, in host memory and on disk | +| `test_sites.py` | reading the cores one site at a time, with prefetching | | `test_logging.py` | the package leaves the host's logging configuration alone | | `test_gpu_backend.py` | the CuPy path; skipped without CuPy or a GPU | @@ -70,6 +75,19 @@ is the behaviour under test, assert it with `pytest.warns`. The exact fallback f networks with fewer than three sites logs a warning rather than raising one, so tests of that path assert on the log instead, as below. +## Exercising the out-of-core paths + +The batched contractions and the host and disk tiers of the sweep only engage when +the budgets are tight, so tests force them with tiny budgets: on the CPU, +`Resources(host_memory="1MB", scratch_dir=tmp_path)` gives small batches and spills +the environments of a depth-4 stack with bonds of 4 to `tmp_path` (see +`tests/test_stack.py`). Compare the dense operator of the result with that of a +default run, not the cores: batching changes the rounding, and with it the cores of +an ill-conditioned sketch, but not the operator they represent. The planner is a +pure function, so `tests/test_plan.py` checks batch sizes and tiers from shapes +alone, and `tests/test_gpu_backend.py` derives GPU budgets from a plan made with +`make_plan` to reach every tier. + ## Logging in tests `log_cli` is on, so log records show up live while the tests run. To assert on diff --git a/docs/content/docs/features/gpu.mdx b/docs/content/docs/features/gpu.mdx index 1a4c751..47c7526 100644 --- a/docs/content/docs/features/gpu.mdx +++ b/docs/content/docs/features/gpu.mdx @@ -8,7 +8,7 @@ through [CuPy](https://cupy.dev/). Both NVIDIA (CUDA) and AMD (ROCm) GPUs are supported through optional extras: ```bash -uv pip install "src_method[gpu-nvidia]" # CUDA 12.x +uv pip install "src_method[gpu-nvidia]" # CUDA 13.x, driver >= 580 uv pip install "src_method[gpu-rocm]" # ROCm ``` diff --git a/docs/content/docs/features/index.mdx b/docs/content/docs/features/index.mdx index f5efb3c..b22522c 100644 --- a/docs/content/docs/features/index.mdx +++ b/docs/content/docs/features/index.mdx @@ -1,6 +1,6 @@ --- title: Features -description: Stacks, adaptive truncation, precision, GPUs and logging. +description: Stacks, adaptive truncation, precision, GPUs, large problems and logging. --- @@ -8,5 +8,6 @@ description: Stacks, adaptive truncation, precision, GPUs and logging. + diff --git a/docs/content/docs/features/large-problems.mdx b/docs/content/docs/features/large-problems.mdx new file mode 100644 index 0000000..3cc4bd7 --- /dev/null +++ b/docs/content/docs/features/large-problems.mdx @@ -0,0 +1,92 @@ +--- +title: Large problems +description: Run stacks that fit neither on the GPU nor in host memory. +--- + +A single SRC sweep keeps, besides the input and output trains, one sketched +environment per site. For stacks with a large bond, such as `N . V . M . U` with a +bond of several thousand in `M`, these environments and the intermediates of the +contractions no longer fit on a GPU, and often not in host memory either. +`src_method` then plans the sweep to the memory at hand: + +- the input cores are read one site at a time, so a train can live on disk; +- every contraction runs in batches of sketch columns (or rows), sized to the GPU + budget; +- each environment is kept on the GPU, in host memory or in a scratch directory, + the most recent ones on the fastest tier. + +None of this changes the result beyond floating-point rounding: the random draws +and the mathematics are those of a sweep that fits in memory. + +## Budgets + +The budgets are set with `Resources`, passed to `src`, `apply` or `compress`: + +```python notest +from src_method import Resources, src + +out = src( + N, V, M, U, + chi_out=2000, + dtype=np.complex128, + device="gpu", + resources=Resources(gpu_memory="36GB", scratch_dir="/local/scratch"), +) +``` + +Every field left unset is detected when the call starts: + +| Field | Default | +|---|---| +| `gpu_memory` | free device memory, plus the free bytes of CuPy's pool, minus `max(10%, 1 GiB)` | +| `host_memory` | `MemAvailable` from `/proc/meminfo`, minus 10% | +| `scratch_dir` | `tempfile.gettempdir()`, which honours `$TMPDIR` | + +Sizes are byte counts or strings: `"36GB"` is `36 * 10**9` bytes, `"36GiB"` is +`36 * 2**30`. On the CPU there is a single budget, `host_memory`. Host memory is +measured when the call starts, so inputs already held in memory are not counted +twice. + +The sweep is planned to fit the GPU budget, counting the bytes its arrays use. +CuPy's memory pool also holds blocks that are split and only partly in use, so +during the sweep the pool is capped at the GPU budget plus the margin of +`max(10%, 1 GiB)` of the device. That margin is the one detection leaves free, so +with a detected budget the cap is the free device memory. An estimate that falls +short by more than the margin fails at once rather than when some other +allocation does; the previous limit is restored afterwards. If a single site does not fit the +budget even with batches of one column, `src` raises `MemoryError` naming the site. + +## Inputs on disk + +A train is any sequence of per-site array-likes with `shape`, `dtype`, `ndim` and +`np.asarray` support (the `SiteLike` protocol, exported for type annotations): +NumPy arrays, `np.memmap`, zarr or HDF5 datasets. Each core +is read when the sweep reaches its site, once per pass, and a background thread +reads the next site while the current one is computed. Raw `.npy` files opened with +`np.load(path, mmap_mode="r")` are the fastest option; compressed formats may be +limited by decompression. + +## The scratch directory + +Environments that fit in neither budget are written to a per-process directory, +`/src--/`, one file per site. Use node-local disk: at the +reference size of 50 sites, `D_M = 4000` and `chi_out = 2000` in complex128, up to +410 GB are written and read back once. The directory is removed when the call +returns, also after an error; a process killed with `SIGKILL` leaves it behind, and +its name identifies the process. + +## Reading the plan + +With `DEBUG` logging on for the `src_method` logger (see +[Logging](/features/logging)), every sweep logs its plan as `SRC plan: ...`: + +- `prefetch`: whether the next site is read ahead; +- `device peak`, `host peak`, `disk`: the planned peaks, in bytes; +- `scratch`: the scratch directory, when anything spills to disk; +- `tiers`: where the environment of each site lives (`device`, `host` or `disk`); +- `batches`: for each site, the batch of the environment, sketch and projection + steps. + +It then logs the time of each pass and, at the end, `SRC stalls`: the seconds +the sweep waited for input cores and for environments. If they are a large +fraction of the run, the disk is too slow for the compute. diff --git a/docs/content/docs/features/meta.json b/docs/content/docs/features/meta.json index b5f5195..272b89d 100644 --- a/docs/content/docs/features/meta.json +++ b/docs/content/docs/features/meta.json @@ -1,4 +1,4 @@ { "title": "Features", - "pages": ["stacks", "cutoff", "precision", "gpu", "logging"] + "pages": ["stacks", "cutoff", "precision", "gpu", "large-problems", "logging"] } diff --git a/docs/content/docs/getting-started.mdx b/docs/content/docs/getting-started.mdx index f1a02ac..f8770af 100644 --- a/docs/content/docs/getting-started.mdx +++ b/docs/content/docs/getting-started.mdx @@ -9,7 +9,7 @@ description: Install src_method and compress your first product. # CPU only (default) uv pip install src_method -# With NVIDIA GPU support (CUDA 12.x) +# With NVIDIA GPU support (CUDA 13.x, driver >= 580) uv pip install "src_method[gpu-nvidia]" # With AMD GPU support (ROCm) diff --git a/pyproject.toml b/pyproject.toml index b4f80a3..7bf32e9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,16 +36,8 @@ dependencies = ["numpy>=2.0,<3", "opt_einsum>=3.4.0"] [project.optional-dependencies] -gpu-nvidia = [ - "cupy-cuda12x>=13", - "nvidia-cuda-nvrtc-cu12", - "nvidia-cuda-runtime-cu12", - "nvidia-cublas-cu12", - "nvidia-cusolver-cu12", - "nvidia-curand-cu12", - "nvidia-cusparse-cu12", - "nvidia-nvjitlink-cu12", -] +# The ctk extra pulls the CUDA 13 libraries from NVIDIA's cuda-toolkit wheels. +gpu-nvidia = ["cupy-cuda13x[ctk]>=14"] gpu-rocm = ["cupy>=14.1.1"] @@ -275,8 +267,8 @@ required-imports = ["from __future__ import annotations"] python-version = "3.11" [tool.ty.analysis] -# Optional GPU dependency, imported lazily in `get_xp`. -allowed-unresolved-imports = ["cupy"] +# Optional GPU dependency, imported lazily in `utils._backend`. +allowed-unresolved-imports = ["cupy", "cupyx"] [tool.ty.src] include = ["src/src_method", "tests"] diff --git a/src/src_method/__init__.py b/src/src_method/__init__.py index aeb6679..2f57c9f 100644 --- a/src/src_method/__init__.py +++ b/src/src_method/__init__.py @@ -7,6 +7,8 @@ import logging +from ._plan import Resources +from ._tensor_train import SiteLike from ._version import version as __version__ from ._version import version_tuple as __version_tuple__ from .apply import apply @@ -16,4 +18,12 @@ # Stay silent unless the application configures logging. logging.getLogger(__name__).addHandler(logging.NullHandler()) -__all__ = ["__version__", "__version_tuple__", "apply", "compress", "src"] +__all__ = [ + "Resources", + "SiteLike", + "__version__", + "__version_tuple__", + "apply", + "compress", + "src", +] diff --git a/src/src_method/_kernels.py b/src/src_method/_kernels.py new file mode 100644 index 0000000..47190d4 --- /dev/null +++ b/src/src_method/_kernels.py @@ -0,0 +1,174 @@ +"""Site kernels of the SRC sweep and the memory they need. + +Every kernel is one of the contractions of the sweep restricted to a batch: a slice +of the sketch index for the environments and the sketch, a slice of the rows of +the new projected environment otherwise. Batching never changes the result beyond +rounding, because the sketch index and those rows are free indices of their +contractions. +""" + +from __future__ import annotations + +from functools import cache +from itertools import count +from math import prod +from typing import TYPE_CHECKING, NamedTuple + +import opt_einsum as oe +from opt_einsum import contract_expression, get_symbol + +if TYPE_CHECKING: + from collections.abc import Sequence + + from numpy.typing import NDArray + from opt_einsum.contract import ContractExpression + +Shape = tuple[int, ...] + + +class Equations(NamedTuple): + """The einsum equations of one sweep, for a fixed stack depth.""" + + ltr: str + rtl_m: str + rtl_s: str + first: str + + +@cache +def equations(depth: int) -> Equations: + """Generate the sweep equations for a stack of ``depth`` layers. + + Layer ``i`` at a site carries ``(a_i, b_i, x_i, x_{i+1})``: left and right + bonds, then its upper and lower physical legs, so that consecutive layers share + ``x``. The output legs are ``x_0`` (up) and ``x_depth`` (down). + """ + symbols = map(get_symbol, count()) + sketch, eta_right, eta_left = next(symbols), next(symbols), next(symbols) + left = "".join(next(symbols) for _ in range(depth)) + right = "".join(next(symbols) for _ in range(depth)) + phys = [next(symbols) for _ in range(depth + 1)] + up, down = phys[0], phys[-1] + layers = ",".join( + f"{left[i]}{right[i]}{phys[i]}{phys[i + 1]}" for i in range(depth) + ) + return Equations( + ltr=f"{sketch}{left},{sketch}{up}{down},{layers}->{sketch}{right}", + rtl_m=f"{sketch}{left},{layers},{eta_right}{right}->{eta_right}{up}{down}{sketch}", + rtl_s=f"{eta_left}{eta_right}{up}{down},{layers},{eta_right}{right}->{eta_left}{left}", + first=f"{layers},{eta_right}{right}->{left}{eta_right}{up}{down}", + ) + + +class _Contractions: + """Compiled contractions for one sweep, keyed on equation and operand shapes. + + Uniform bulk sites share one entry, so the path is planned once rather than at + every site. Kept per call: jagged bonds and short last batches add entries that + are not worth keeping. + """ + + def __init__(self) -> None: + self._compiled: dict[tuple[str, tuple[Shape, ...]], ContractExpression] = {} + + def __call__(self, eq: str, *operands: NDArray) -> NDArray: + shapes = tuple(op.shape for op in operands) + expr = self._compiled.get((eq, shapes)) + if expr is None: + expr = self._compiled[eq, shapes] = contract_expression(eq, *shapes) + return expr(*operands) + + +class SiteKernels: + """The four contractions of the sweep, each applied to one batch. + + Args: + depth: The number of layers of the stack. + """ + + def __init__(self, depth: int) -> None: + self.eqs = equations(depth) + self._contract = _Contractions() + + def env(self, env: NDArray, omega: NDArray, cores: Sequence[NDArray]) -> NDArray: + """Advance a batch of sketch columns of the environment by one site. + + Args: + env: Columns ``lo:hi`` of ``C_j``, shape ``(b, *left_bonds)``. + omega: The same columns of the site's Gaussian tensor, ``(b, up, down)``. + cores: The padded cores of the site. + + Returns: + Columns ``lo:hi`` of ``C_{j+1}``, shape ``(b, *right_bonds)``. + """ + return self._contract(self.eqs.ltr, env, omega, *cores) + + def sketch(self, env: NDArray, cores: Sequence[NDArray], proj: NDArray) -> NDArray: + """Sketch a batch of columns of the site's running core. + + Args: + env: Columns ``lo:hi`` of ``C_j``, shape ``(b, *left_bonds)``. + cores: The padded cores of the site. + proj: The projected environment ``S``, ``(eta, *right_bonds)``. + + Returns: + Columns ``lo:hi`` of the sketch, shape ``(eta, up, down, b)``. + """ + return self._contract(self.eqs.rtl_m, env, *cores, proj) + + def project(self, eta: NDArray, cores: Sequence[NDArray], proj: NDArray) -> NDArray: + """Project a batch of rows of the new projected environment. + + Args: + eta: Rows ``lo:hi`` of the conjugated output core, ``(b, eta, up, down)``. + cores: The padded cores of the site. + proj: The projected environment ``S``, ``(eta, *right_bonds)``. + + Returns: + Rows ``lo:hi`` of the new ``S``, shape ``(b, *left_bonds)``. + """ + return self._contract(self.eqs.rtl_s, eta, *cores, proj) + + def first(self, cores: Sequence[NDArray], proj: NDArray) -> NDArray: + """Contract the first site with a batch of rows of ``S``. + + Args: + cores: The padded cores of the first site. + proj: Rows ``lo:hi`` of ``S``, ``(b, *right_bonds)``. + + Returns: + The output core for those rows, ``(*left_bonds, b, up, down)``. + """ + return self._contract(self.eqs.first, *cores, proj) + + +@cache +def peak_elements(eq: str, shapes: tuple[Shape, ...]) -> int: + """Estimate the peak elements a contraction allocates beyond its inputs. + + Walks the path `opt_einsum` picks for these shapes, the one `SiteKernels` runs. + Each pairwise step holds the intermediates still alive, its output and a + possible contiguous copy of both operands (``tensordot`` transposes them), so + the estimate errs on the high side. The final output is included, the inputs + are not. + + Args: + eq: The einsum equation. + shapes: The operand shapes. + + Returns: + The peak number of elements. + """ + _, info = oe.contract_path(eq, *shapes, shapes=True) + sizes = info.size_dict + # (elements, is_intermediate) for every operand still to be contracted. + operands = [(prod(shape), False) for shape in shapes] + peak = 0 + for step in info.contraction_list: + positions, einsum_str = step[0], step[2] + live = sum(n for n, is_tmp in operands if is_tmp) + popped = [operands.pop(i) for i in positions] + out = prod(sizes[c] for c in einsum_str.split("->")[1]) + peak = max(peak, live + out + sum(n for n, _ in popped)) + operands.append((out, True)) + return peak diff --git a/src/src_method/_plan.py b/src/src_method/_plan.py new file mode 100644 index 0000000..aa44299 --- /dev/null +++ b/src/src_method/_plan.py @@ -0,0 +1,536 @@ +"""Memory planning for the SRC sweep: budgets, batch sizes and environment tiers. + +`make_plan` is a pure function of the core shapes, the sketch size, the dtype and +the budgets, so the plan of a run can be inspected, and tested, without loading any +data or touching a GPU. +""" + +from __future__ import annotations + +import re +import shutil +import tempfile +from dataclasses import dataclass +from math import prod +from pathlib import Path +from typing import TYPE_CHECKING, Literal + +import numpy as np + +from ._kernels import equations, peak_elements +from .utils import device_memory, host_memory_available, is_host + +if TYPE_CHECKING: + import os + from collections.abc import Callable, Sequence + from types import ModuleType + +Tier = Literal["device", "host", "disk"] +Shape = tuple[int, ...] + +# Device memory outside the plan: cuBLAS/cuSOLVER workspaces and pool fragmentation. +# Detection keeps it free, and CuPy's pool may grow into it during the sweep. +GPU_MARGIN_FRACTION = 0.10 +GPU_MARGIN_MIN = 2**30 +HOST_MARGIN_FRACTION = 0.10 +DISK_MARGIN_FRACTION = 0.05 +# Batch size used to decide how much of the device can hold environments. +PREFERRED_BATCH = 512 +# Batches are rounded down to a multiple of this for efficient GEMMs. +GEMM_MULTIPLE = 32 + +_UNITS = { + "B": 1, + "KB": 10**3, + "MB": 10**6, + "GB": 10**9, + "TB": 10**12, + "KIB": 2**10, + "MIB": 2**20, + "GIB": 2**30, + "TIB": 2**40, +} +_SIZE = re.compile(r"^\s*(\d+(?:\.\d+)?)\s*([KMGT]i?B|B)?\s*$", re.IGNORECASE) + + +def parse_size(value: int | str) -> int: + """Convert a byte count or a size string to bytes. + + Args: + value: A non-negative integer, or a string such as ``"36GB"`` (decimal, + ``36 * 10**9``) or ``"36GiB"`` (binary, ``36 * 2**30``). + + Returns: + The size in bytes. + + Raises: + TypeError: If the value is neither an integer nor a string. + ValueError: If the value is negative or not a recognised size. + """ + if isinstance(value, bool) or not isinstance(value, int | str): + msg = f"Expected a byte count or a size string such as '36GB', got {value!r}." + raise TypeError(msg) + if isinstance(value, int): + if value < 0: + msg = f"Expected a non-negative byte count, got {value}." + raise ValueError(msg) + return value + match = _SIZE.match(value) + if match is None: + msg = f"Expected a size string such as '36GB' or '36GiB', got {value!r}." + raise ValueError(msg) + number, unit = match.groups() + return int(float(number) * _UNITS[(unit or "B").upper()]) + + +@dataclass(frozen=True) +class Resources: + """Memory budgets and scratch space for one `src` call. + + Every field left as ``None`` is detected when the call starts. + + Attributes: + gpu_memory: Device memory the sweep plans to use, as a byte count or a + size string (``"36GB"``, ``"36GiB"``). Defaults to the free device + memory minus a margin of ``max(10%, 1 GiB)`` of the device. CuPy's pool + is capped at the budget plus that margin, which absorbs fragmentation. + Ignored on the CPU. + host_memory: Host memory the sweep may use. Defaults to ``MemAvailable`` + minus 10%. On the CPU it covers the working set as well. + scratch_dir: Directory for environments that fit in neither budget, + ideally on node-local disk. Defaults to ``tempfile.gettempdir()``, + which honours ``$TMPDIR``. + """ + + gpu_memory: int | str | None = None + host_memory: int | str | None = None + scratch_dir: str | os.PathLike[str] | None = None + + def __post_init__(self) -> None: + """Validate the explicit budgets. + + Raises: + TypeError: If a budget is neither an integer nor a string. + ValueError: If a budget is not a recognised size. + """ + for value in (self.gpu_memory, self.host_memory): + if value is not None: + parse_size(value) + + +@dataclass(frozen=True) +class Budgets: + """The resolved budgets of one sweep, in bytes. + + Attributes: + device: Bytes for the working set and the device tier. + host: Bytes for the output, the staging buffers and the host tier. + disk: Bytes free in ``scratch_dir``. + scratch_dir: Where the disk tier lives. + unified: Whether device and host memory are the same (the CPU backend). + device_cap: The cap on CuPy's pool during the sweep: ``device`` plus the + margin for workspaces and fragmentation, which the plan leaves out. + ``None`` leaves the pool uncapped, as on the CPU backend. + """ + + device: int + host: int + disk: int + scratch_dir: Path + unified: bool + device_cap: int | None = None + + +def resolve_budgets(resources: Resources | None, xp: ModuleType) -> 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``). + + Returns: + The budgets. + """ + resources = resources or Resources() + scratch = ( + Path(resources.scratch_dir) + if resources.scratch_dir is not None + else Path(tempfile.gettempdir()) + ) + if resources.host_memory is not None: + host = parse_size(resources.host_memory) + else: + host = int(host_memory_available() * (1 - HOST_MARGIN_FRACTION)) + unified = is_host(xp) + cap = None + if unified: + device = host + else: + available, total = device_memory(xp) + margin = max(int(GPU_MARGIN_FRACTION * total), GPU_MARGIN_MIN) + if resources.gpu_memory is not None: + device = max(parse_size(resources.gpu_memory), 0) + else: + device = max(available - margin, 0) + # The plan counts the bytes in use, but the pool limit applies to every + # block the pool holds, including split blocks that are partly free. + cap = device + margin + free = shutil.disk_usage(_existing_parent(scratch)).free + disk = int(free * (1 - DISK_MARGIN_FRACTION)) + return Budgets(device, max(host, 0), disk, scratch, unified=unified, device_cap=cap) + + +def _existing_parent(path: Path) -> Path: + """Return ``path`` or its nearest existing ancestor.""" + path = path.absolute() + while not path.exists(): + path = path.parent + return path + + +@dataclass(frozen=True) +class SitePlan: + """How one site is processed. + + Attributes: + env_batch: Sketch columns per environment step (0 at the last site). + sketch_batch: Sketch columns per sketch step (0 at the first site). + project_batch: Rows per projection step, or per first-site step at site 0. + tier: Where ``C_j`` is kept; site 0 stores no environment. + """ + + env_batch: int + sketch_batch: int + project_batch: int + tier: Tier + + +@dataclass(frozen=True) +class Plan: + """The memory plan of one sweep. + + Attributes: + sites: One entry per site. + prefetch: How many sites ahead the cores are loaded (0 or 1). + device_peak: Estimated peak device bytes. + host_peak: Estimated peak host bytes, beyond the inputs. + disk_bytes: Bytes spilled to the scratch directory. + """ + + sites: tuple[SitePlan, ...] + prefetch: int + device_peak: int + host_peak: int + disk_bytes: int + + +class _Site: + """The memory model of one site, in bytes.""" + + def __init__( + self, + j: int, + shapes: tuple[Shape, ...], + core_bytes: int, + *, + chi: int, + itemsize: int, + n_sites: int, + ) -> None: + self.j, self.shapes, self.core_bytes, self.chi, self.e = ( + j, + shapes, + core_bytes, + chi, + itemsize, + ) + self.eqs = equations(len(shapes)) + self.left = tuple(s[0] for s in shapes) + self.right = tuple(s[1] for s in shapes) + self.a, self.b = prod(self.left), prod(self.right) + self.up, self.down = shapes[0][2], shapes[-1][3] + self.p = self.up * self.down + # Rows of the projected environment S that enters site j right-to-left. + self.eta = 1 if j == n_sites - 1 else chi + self.env_bytes = chi * self.a * itemsize + + def env(self, b: int, prefetch: int, *, staged_in: bool, staged_out: bool) -> int: + """Bytes of one environment step on ``b`` columns.""" + e, chi = self.e, self.chi + fixed = self.core_bytes * (1 + prefetch) + chi * self.p * e + slices = (1 + staged_in) * b * self.a * e + staged_out * b * self.b * e + peak = peak_elements( + self.eqs.ltr, ((b, *self.left), (b, self.up, self.down), *self.shapes) + ) + return fixed + slices + peak * e + + def sketch(self, b: int, prefetch: int, *, staged: bool) -> int: + """Bytes of one sketch step on ``b`` columns.""" + e, chi = self.e, self.chi + fixed = ( + self.core_bytes * (1 + prefetch) + + self.eta * self.b * e + + self.eta * self.p * chi * e + ) + slices = (1 + staged) * b * self.a * e + peak = ( + peak_elements( + self.eqs.rtl_m, + ((b, *self.left), *self.shapes, (self.eta, *self.right)), + ) + * e + ) + return fixed + slices + peak + + def qr(self, prefetch: int) -> int: + """Bytes of the QR of the sketch: the sketch, ``Q`` and a workspace.""" + e = self.e + return ( + self.core_bytes * (1 + prefetch) + + self.eta * self.b * e + + 3 * self.eta * self.p * self.chi * e + ) + + def project(self, b: int, prefetch: int) -> int: + """Bytes of one projection step on ``b`` rows.""" + e, chi = self.e, self.chi + fixed = ( + self.core_bytes * (1 + prefetch) + + self.eta * self.b * e + + self.eta * self.p * chi * e + + chi * self.a * e + ) + slices = b * self.eta * self.p * e + peak = ( + peak_elements( + self.eqs.rtl_s, + ( + (b, self.eta, self.up, self.down), + *self.shapes, + (self.eta, *self.right), + ), + ) + * e + ) + return fixed + slices + peak + + def first(self, b: int, prefetch: int) -> int: + """Bytes of one first-site step on ``b`` rows of ``S``.""" + e, chi = self.e, self.chi + fixed = ( + self.core_bytes * (1 + prefetch) + + self.eta * self.b * e + + self.a * chi * self.p * e + ) + peak = peak_elements(self.eqs.first, (*self.shapes, (b, *self.right))) * e + return fixed + peak + + +def _largest_batch(cost: Callable[[int], int], limit: int, avail: int) -> int: + """Return the largest batch in ``[1, limit]`` whose cost fits, or 0. + + Binary search on the evaluated costs; the result is always a batch whose cost + was checked. Batches of at least `GEMM_MULTIPLE` are rounded down to a multiple + of it when that still fits. + """ + if cost(1) > avail: + return 0 + lo, hi = 1, limit + while lo < hi: + mid = (lo + hi + 1) // 2 + if cost(mid) <= avail: + lo = mid + else: + hi = mid - 1 + rounded = lo - lo % GEMM_MULTIPLE + if lo >= GEMM_MULTIPLE and rounded != lo and cost(rounded) <= avail: + return rounded + return lo + + +class _Planner: + """Plans one sweep for a fixed prefetch depth.""" + + def __init__(self, sites: Sequence[_Site], budgets: Budgets, prefetch: int) -> None: + self.sites, self.budgets, self.prefetch = sites, budgets, prefetch + self.n = len(sites) + self.chi = sites[0].chi + self.out_total = sum(self.chi * s.p * s.eta * s.e for s in sites) + # Pinned ring of the site source, plus the host copy being read. + self.site_ring = (prefetch + 2) * max(s.core_bytes for s in sites) + # Two batches in and two out; full environments until batches are known. + self.env_staging = 4 * max(s.env_bytes for s in sites[1:]) + + def _costs(self, j: int, tiers: Sequence[Tier]) -> dict[str, Callable[[int], int]]: + site, pf = self.sites[j], self.prefetch + costs: dict[str, Callable[[int], int]] = {} + if j < self.n - 1: + staged_in = j > 0 and tiers[j] != "device" + staged_out = tiers[j + 1] != "device" + costs["env"] = lambda b: site.env( + b, pf, staged_in=staged_in, staged_out=staged_out + ) + if j > 0: + staged = tiers[j] != "device" + costs["sketch"] = lambda b: site.sketch(b, pf, staged=staged) + costs["project"] = lambda b: site.project(b, pf) + else: + costs["first"] = lambda b: site.first(b, pf) + return costs + + def _batches(self, tiers: Sequence[Tier], avail: int) -> list[dict[str, int]]: + batches = [] + for j, site in enumerate(self.sites): + if j > 0 and site.qr(self.prefetch) > avail: + _fail(j, "QR", site.qr(self.prefetch), avail) + chosen = {} + for kernel, cost in self._costs(j, tiers).items(): + b = _largest_batch(cost, self.chi, avail) + if b == 0: + _fail(j, kernel, cost(1), avail) + chosen[kernel] = b + batches.append(chosen) + return batches + + def _peak(self, tiers: Sequence[Tier], batches: Sequence[dict[str, int]]) -> int: + return max( + cost(batches[j][kernel]) + for j in range(self.n) + for kernel, cost in self._costs(j, tiers).items() + ) + + def _reserved(self, tiers: Sequence[Tier]) -> int: + """Device bytes held outside the kernels. + + These are the device-tier environments and, on the CPU backend, where the + device is the host, the output and the staging buffers as well. + """ + resident = sum( + s.env_bytes + for j, s in enumerate(self.sites) + if j > 0 and tiers[j] == "device" + ) + if self.budgets.unified: + spills = any(tier == "disk" for tier in tiers[1:]) + resident += self.out_total + self.site_ring + spills * self.env_staging + return resident + + def _tiers(self, device_left: int, host_left: int) -> list[Tier]: + """Assign the newest environments to the fastest tier that holds them.""" + tiers: list[Tier] = ["device"] * self.n + level: Tier = "device" + for j in range(self.n - 1, 0, -1): + need = self.sites[j].env_bytes + if level == "device" and device_left >= need: + device_left -= need + continue + if level == "device": + level = "disk" if self.budgets.unified else "host" + if level == "host" and host_left >= need: + tiers[j], host_left = "host", host_left - need + continue + tiers[j] = level = "disk" + return tiers + + def plan(self) -> Plan: + budgets, n = self.budgets, self.n + staged: list[Tier] = ["device"] + ["disk"] * (n - 1) + base = budgets.device - self._reserved(staged) + batches = self._batches(staged, base) + # Staged batches never grow in the final pass, so they bound the staging. + self.env_staging = 4 * max( + max(batches[j].values()) * s.a * s.e for j, s in enumerate(self.sites) + ) + + preferred = [ + {kernel: min(b, PREFERRED_BATCH) for kernel, b in chosen.items()} + for chosen in batches + ] + host_left = budgets.host - self.out_total - self.site_ring - self.env_staging + tiers = self._tiers(base - self._peak(staged, preferred), host_left) + + disk_bytes = sum( + s.env_bytes for j, s in enumerate(self.sites) if tiers[j] == "disk" + ) + if disk_bytes > budgets.disk: + msg = ( + f"The environments need {disk_bytes} bytes on disk, but only " + f"{budgets.disk} bytes are free in {budgets.scratch_dir}." + ) + raise MemoryError(msg) + + reserved = self._reserved(tiers) + batches = self._batches(tiers, budgets.device - reserved) + host_env = sum( + s.env_bytes for j, s in enumerate(self.sites) if tiers[j] == "host" + ) + host_peak = self.out_total + self.site_ring + if disk_bytes or host_env: + host_peak += host_env + self.env_staging + sites = tuple( + SitePlan( + env_batch=chosen.get("env", 0), + sketch_batch=chosen.get("sketch", 0), + project_batch=chosen.get("project", chosen.get("first", 0)), + tier=tiers[j], + ) + for j, chosen in enumerate(batches) + ) + return Plan( + sites, + self.prefetch, + device_peak=reserved + self._peak(tiers, batches), + host_peak=host_peak, + disk_bytes=disk_bytes, + ) + + +def _fail(j: int, kernel: str, need: int, avail: int) -> None: + msg = ( + f"Site {j}: the {kernel} step needs {need} bytes with a batch of one, but " + f"only {avail} bytes of the device budget are available. The site working " + "set exceeds the budget." + ) + raise MemoryError(msg) + + +def make_plan( + site_shapes: Sequence[tuple[Shape, ...]], + site_bytes: Sequence[int], + chi_out: int, + dtype: type | np.dtype, + budgets: Budgets, +) -> Plan: + """Plan the batches and the environment tiers of one sweep. + + Args: + site_shapes: For every site, the padded ``(l, r, u, d)`` shape of each layer. + site_bytes: For every site, the bytes of its cores. + chi_out: The sketch size. + dtype: The data type of the computation. + budgets: The resolved budgets. + + Returns: + The plan. Prefetching is dropped before giving up. + + Raises: + MemoryError: If a site does not fit the device budget even with batches of + one and no prefetching, or the environments do not fit on disk. + """ + itemsize = np.dtype(dtype).itemsize + n_sites = len(site_shapes) + sites = [ + _Site( + j, + tuple(shapes), + core_bytes, + chi=chi_out, + itemsize=itemsize, + n_sites=n_sites, + ) + for j, (shapes, core_bytes) in enumerate(zip(site_shapes, site_bytes)) + ] + try: + return _Planner(sites, budgets, prefetch=1).plan() + except MemoryError: + return _Planner(sites, budgets, prefetch=0).plan() diff --git a/src/src_method/_sites.py b/src/src_method/_sites.py new file mode 100644 index 0000000..d59544d --- /dev/null +++ b/src/src_method/_sites.py @@ -0,0 +1,193 @@ +"""Lazy, per-site access to the cores of a stack. + +A core is read from its source (a NumPy array, an `np.memmap`, a zarr or HDF5 +dataset) only when the sweep reaches its site, and a background thread can read the +next site while the current one is computed. Shapes come from ``.shape`` alone, so +planning never reads any data. +""" + +from __future__ import annotations + +import queue +from concurrent.futures import Future, ThreadPoolExecutor +from math import prod +from time import perf_counter +from typing import TYPE_CHECKING, Any, Self + +import numpy as np + +from ._tensor_train import known_kind, pad_site, padded_shape +from .utils import NullEvent, current_stream, is_host, pinned_empty, to_device_async + +if TYPE_CHECKING: + from collections.abc import Sequence + from types import ModuleType, TracebackType + + from numpy.typing import NDArray + + from ._tensor_train import Site + +# Byte alignment of each core inside a page-locked staging buffer. +_ALIGN = 256 + + +def padded_shapes( + layers: Sequence[Sequence[Site]], +) -> list[tuple[tuple[int, ...], ...]]: + """Return, for every site, the padded ``(l, r, u, d)`` shape of each layer. + + Args: + layers: A stack in ket form. + + Returns: + One tuple of shapes per site, without reading any data. + """ + kinds = [known_kind(layer) for layer in layers] + last = len(layers[0]) - 1 + return [ + tuple( + padded_shape(tuple(layer[j].shape), kind, j, last) + for layer, kind in zip(layers, kinds) + ) + for j in range(last + 1) + ] + + +def site_bytes(layers: Sequence[Sequence[Site]]) -> list[int]: + """Return, for every site, the bytes of its cores over all layers.""" + return [ + sum( + prod(layer[j].shape) * np.dtype(layer[j].dtype).itemsize for layer in layers + ) + for j in range(len(layers[0])) + ] + + +def _aligned(n_bytes: int) -> int: + return -(-n_bytes // _ALIGN) * _ALIGN + + +class SiteSource: + """Padded device cores of a stack, one site at a time. + + Reading from the source and staging into page-locked memory happen on a + background thread; the host-to-device copy is issued on the current stream when + the site is requested, so it is ordered with the kernels that use it. + + Use as a context manager, so that the loader thread stops. + + 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``). + depth: How many sites `prefetch` may load ahead (0 disables the thread). + """ + + def __init__( + self, layers: Sequence[Sequence[Site]], xp: ModuleType, *, depth: int + ) -> None: + self._layers = layers + self._kinds = [known_kind(layer) for layer in layers] + self._last = len(layers[0]) - 1 + self._xp = xp + self._pending: dict[int, Future[tuple[list[Any], Any]]] = {} + self._executor = ( + ThreadPoolExecutor(1, thread_name_prefix="src-site-loader") + if depth > 0 + else None + ) + self._buffers: queue.Queue[tuple[Any, Any]] = queue.Queue() + self.stall_seconds = 0.0 + if not is_host(xp): + n_bytes = max( + sum( + _aligned(prod(layer[j].shape) * np.dtype(layer[j].dtype).itemsize) + for layer in layers + ) + for j in range(self._last + 1) + ) + for _ in range(depth + 1): + self._buffers.put((pinned_empty(n_bytes, xp), NullEvent())) + + def __len__(self) -> int: + """Return the number of sites.""" + return self._last + 1 + + def __enter__(self) -> Self: + """Return the source.""" + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + """Stop the loader thread.""" + if self._executor is not None: + self._executor.shutdown(wait=True, cancel_futures=True) + + def prefetch(self, j: int) -> None: + """Start reading site ``j`` in the background, if there is such a site.""" + if self._executor is None or not 0 <= j <= self._last or j in self._pending: + return + self._pending[j] = self._executor.submit(self._stage, j) + + def __getitem__(self, j: int) -> tuple[NDArray, ...]: + """Return the padded cores of site ``j`` on the device.""" + future = self._pending.pop(j, None) + if future is None: + staged, buffer = self._stage(j) + else: + start = perf_counter() + staged, buffer = future.result() + self.stall_seconds += perf_counter() - start + if is_host(self._xp): + return tuple(staged) + stream = current_stream(self._xp) + cores = [ + core + if isinstance(core, self._xp.ndarray) + else to_device_async(core, self._xp, stream) + for core in staged + ] + # The buffer is refilled only once these copies are done. + self._buffers.put((buffer, stream.record())) + return tuple( + pad_site(core, kind, j, self._last) + for core, kind in zip(cores, self._kinds) + ) + + def _stage(self, j: int) -> tuple[list[Any], Any]: + """Read site ``j`` into memory. + + Returns padded arrays on the host backend, and on the GPU unpadded + page-locked copies (or the device arrays given as input) with their buffer. + """ + if is_host(self._xp): + padded = [ + pad_site(_read(layer[j]), kind, j, self._last) + for layer, kind in zip(self._layers, self._kinds) + ] + return padded, None + buffer, previous = self._buffers.get() + previous.synchronize() + staged, offset = [], 0 + for layer in self._layers: + core = layer[j] + if not isinstance(core, self._xp.ndarray): + host = np.asarray(core) + view = buffer[offset : offset + host.nbytes] + view = view.view(host.dtype).reshape(host.shape) + np.copyto(view, host) + offset += _aligned(host.nbytes) + core = view + staged.append(core) + return staged, buffer + + +def _read(core: Site) -> 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) + return core diff --git a/src/src_method/_store.py b/src/src_method/_store.py new file mode 100644 index 0000000..1abbf3e --- /dev/null +++ b/src/src_method/_store.py @@ -0,0 +1,347 @@ +"""Where the sweep keeps its environments: on the device, in host memory or on disk. + +The right-to-left pass reads the environments in the reverse order of the +left-to-right pass that writes them, one batch of sketch columns at a time. Off the +device, every batch moves through a small ring of page-locked buffers: a writer +thread drains device-to-host copies into host arrays or files, and a reader thread +fills buffers ahead of use, so that transfers overlap the kernels. +""" + +from __future__ import annotations + +import os +import queue +import shutil +import uuid +import weakref +from collections import deque +from concurrent.futures import Future, ThreadPoolExecutor +from math import prod +from time import perf_counter +from typing import TYPE_CHECKING, Any, Self + +import numpy as np + +from .utils import current_stream, pinned_empty, to_device_async, to_host_async + +if TYPE_CHECKING: + from collections.abc import Iterable, Sequence + from pathlib import Path + from types import ModuleType, TracebackType + + from numpy.typing import NDArray + + from ._plan import Plan + +# Page-locked buffers per direction: one in use, one in flight. +RING_SIZE = 2 +# How often a blocked reader checks whether the store is closing. +_POLL_SECONDS = 0.1 + + +class _Closing(Exception): # noqa: N818 (control flow, not an error) + """Raised in the reader thread when the store closes under it.""" + + +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. + + 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. + """ + + def __init__( + self, + 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) + ) -> None: + self._plan = plan + 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 + self._copy = copy_stream + self._device: dict[int, NDArray] = {} + self._host: dict[int, np.ndarray] = {} + self._fds: dict[int, int] = {} + self._writes: dict[int, list[Future[None]]] = {} + self._ahead: deque[tuple[int, int, int]] = deque() + self._inflight: deque[tuple[int, int, int, Future[tuple[Any, np.ndarray]]]] = ( + deque() + ) + self._closing = False + self.stall_seconds = 0.0 + + self._dir: Path | None = None + if "disk" in self._tiers: + self._dir = scratch_dir / f"src-{os.getpid()}-{uuid.uuid4().hex[:8]}" + self._dir.mkdir(parents=True) + self._finalizer = weakref.finalize( + self, shutil.rmtree, self._dir, ignore_errors=True + ) + + staged = [j for j, tier in enumerate(self._tiers) if j > 0 and tier != "device"] + self._writer = ThreadPoolExecutor(1, thread_name_prefix="src-env-writer") + 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)) + + def _row_bytes(self, j: int) -> int: + return prod(self._shapes[j][1:]) * self._dtype.itemsize + + def _batch_rows(self, j: int) -> int: + sites = self._plan.sites + return max(sites[j - 1].env_batch, sites[j].env_batch, sites[j].sketch_batch) + + def __enter__(self) -> Self: + """Return the store.""" + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + """Close the store; see `close`.""" + self.close(wait=exc is None) + + def close(self, *, wait: bool = True) -> None: + """Stop the threads and remove the scratch files. + + Args: + wait: Wait for the pending writes and re-raise their first error. + """ + try: + if wait: + for futures in self._writes.values(): + for future in futures: + future.result() + finally: + self._closing = True + self._writer.shutdown(wait=True, cancel_futures=True) + self._reader.shutdown(wait=True, cancel_futures=True) + for fd in self._fds.values(): + os.close(fd) + self._fds.clear() + self._device.clear() + self._host.clear() + if self._dir is not None: + self._finalizer() + + # --- writing --- + + def put(self, j: int, lo: int, hi: int, x: NDArray) -> None: + """Store columns ``lo:hi`` of ``C_j``. + + Args: + j: The site. + lo: The first column. + hi: One past the last column. + x: The columns, ``(hi - lo, *left_bonds)``, on the device. + """ + if self._tiers[j] == "device": + if j not in self._device: + self._device[j] = self._xp.empty(self._shapes[j], dtype=self._dtype) + self._device[j][lo:hi] = x + return + self._raise_failed() + self._open(j) + 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()) + to_host_async(x, host, self._copy) + done = self._copy.record() + # x stays referenced by the task until the copy is complete. + future = self._writer.submit( + self._write, j, lo, x, host=host, done=done, buffer=buffer + ) + self._writes.setdefault(j, []).append(future) + + def _open(self, j: int) -> None: + if self._tiers[j] == "host" and j not in self._host: + self._host[j] = np.empty(self._shapes[j], dtype=self._dtype) + elif self._tiers[j] == "disk" and j not in self._fds: + fd = os.open(self._path(j), os.O_RDWR | os.O_CREAT, 0o600) + self._fds[j] = fd + os.ftruncate(fd, prod(self._shapes[j]) * self._dtype.itemsize) + + def _path(self, j: int) -> Path: + assert self._dir is not None # noqa: S101 (the disk tier creates it) + return self._dir / f"env-{j:04d}.bin" + + def _write( + self, + j: int, + lo: int, + x: NDArray, + *, + host: np.ndarray, + done: Any, # noqa: ANN401 (a cupy or null event) + buffer: Any, # noqa: ANN401 (a page-locked buffer) + ) -> None: + try: + done.synchronize() + del x + if self._tiers[j] == "host": + np.copyto(self._host[j][lo : lo + host.shape[0]], host) + return + offset = lo * self._row_bytes(j) + try: + _write_all(self._fds[j], memoryview(buffer[: host.nbytes]), offset) + except OSError as err: + msg = ( + f"Writing environment {j} to {self._path(j)} failed " + f"({err.strerror}); the plan spills {self._plan.disk_bytes} bytes " + f"to {self._dir}." + ) + raise OSError(err.errno, msg) from err + if hasattr(os, "posix_fadvise"): + os.posix_fadvise( + self._fds[j], offset, host.nbytes, os.POSIX_FADV_DONTNEED + ) + finally: + self._out.put(buffer) + + def _raise_failed(self) -> None: + """Re-raise the first error of a finished write, dropping finished writes.""" + for j, futures in self._writes.items(): + pending = [] + for future in futures: + if not future.done(): + pending.append(future) + elif (error := future.exception()) is not None: + raise error + self._writes[j] = pending + + # --- reading --- + + def prefetch(self, j: int, ranges: Iterable[tuple[int, int]]) -> None: + """Announce the batches of ``C_j`` that `get` will be asked for, in order. + + Off the device, the reader starts filling buffers with the first ones. + + Args: + j: The site. + ranges: The ``(lo, hi)`` column ranges, in the order they will be read. + """ + if self._tiers[j] == "device": + return + self._ahead.extend((j, lo, hi) for lo, hi in ranges) + self._schedule() + + def _schedule(self) -> None: + while self._ahead and len(self._inflight) < RING_SIZE: + j, lo, hi = self._ahead.popleft() + writes = list(self._writes.get(j, [])) + future = self._reader.submit(self._read, j, lo, hi, writes) + self._inflight.append((j, lo, hi, future)) + + def _read( + self, j: int, lo: int, hi: int, writes: Sequence[Future[None]] + ) -> tuple[Any, np.ndarray]: + for write in writes: + write.result() + while True: + try: + buffer, event = self._in.get(timeout=_POLL_SECONDS) + break + except queue.Empty: + if self._closing: + raise _Closing from None + if event is not None: + event.synchronize() + shape = (hi - lo, *self._shapes[j][1:]) + n_bytes = prod(shape) * self._dtype.itemsize + host = buffer[:n_bytes].view(self._dtype).reshape(shape) + if self._tiers[j] == "host": + np.copyto(host, self._host[j][lo:hi]) + else: + _read_all( + self._fds[j], memoryview(buffer[:n_bytes]), lo * self._row_bytes(j) + ) + return buffer, host + + def get(self, j: int, lo: int, hi: int) -> NDArray: + """Return columns ``lo:hi`` of ``C_j`` on the device. + + Args: + j: The site. + lo: The first column. + hi: One past the last column. + + Returns: + The columns, ``(hi - lo, *left_bonds)``. Kernels on the current stream + see them once their transfer is done. + + Raises: + RuntimeError: If the batch is not the next one announced by `prefetch`. + """ + if self._tiers[j] == "device": + return self._device[j][lo:hi] + self._raise_failed() + if not self._inflight and not self._ahead: + self.prefetch(j, [(lo, hi)]) + if not self._inflight or self._inflight[0][:3] != (j, lo, hi): + msg = f"Batch {lo}:{hi} of environment {j} was read out of order." + raise RuntimeError(msg) + future = self._inflight.popleft()[3] + start = perf_counter() + buffer, host = future.result() + 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) + self._in.put((buffer, stream.record())) + self._schedule() + return device + + def drop(self, j: int) -> None: + """Release ``C_j``: its memory, or its file.""" + for future in self._writes.pop(j, []): + future.result() + self._device.pop(j, None) + self._host.pop(j, None) + fd = self._fds.pop(j, None) + if fd is not None: + os.close(fd) + self._path(j).unlink() + + +def _write_all(fd: int, data: memoryview, offset: int) -> None: + """Write all of ``data`` at ``offset``; `os.pwrite` may write less.""" + while data: + written = os.pwrite(fd, data, offset) + data, offset = data[written:], offset + written + + +def _read_all(fd: int, data: memoryview, offset: int) -> None: + """Fill ``data`` from ``offset``; `os.preadv` may read less.""" + while data: + read = os.preadv(fd, [data], offset) + if read == 0: + msg = f"Unexpected end of file at offset {offset}." + raise OSError(msg) + data, offset = data[read:], offset + read diff --git a/src/src_method/_sweep.py b/src/src_method/_sweep.py index 3529b5d..c0dac32 100644 --- a/src/src_method/_sweep.py +++ b/src/src_method/_sweep.py @@ -9,91 +9,54 @@ tensors and accumulates the environments ``C``; the sketch index is shared by every site (a Khatri-Rao sketch). The right-to-left sweep builds the output through `truncated_qr` while carrying the projected environment ``S``. + +Every contraction runs in batches sized by `src_method._plan.make_plan` to the +memory budgets. The cores are read one site at a time (`SiteSource`) and the +environments live on the device, in host memory or on disk (`EnvironmentStore`). """ from __future__ import annotations import logging -from functools import cache -from itertools import count -from math import prod from time import perf_counter_ns -from typing import TYPE_CHECKING, NamedTuple - -from opt_einsum import contract_expression, get_symbol - -from ._tensor_train import pad, unpad -from .utils import gaussian_sketch, to_numpy, truncated_qr +from typing import TYPE_CHECKING + +import numpy as np + +from ._kernels import SiteKernels +from ._plan import make_plan, resolve_budgets +from ._sites import SiteSource, padded_shapes, site_bytes +from ._store import EnvironmentStore +from ._tensor_train import unpad +from .utils import ( + device_pool_bytes, + device_pool_limit, + gaussian_sketch, + new_stream, + to_numpy, + truncated_qr, +) if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import Iterator, Sequence from types import ModuleType - import numpy as np from numpy.typing import DTypeLike, NDArray - from opt_einsum.contract import ContractExpression - from ._tensor_train import TrainKind + from ._plan import Plan, Resources + from ._tensor_train import Site, TrainKind logger = logging.getLogger(__name__) -class _Contractions: - """Compiled contractions for one sweep, keyed on equation and operand shapes. - - Uniform bulk sites share one entry, so the path is planned once rather than at - every site. Kept per call: jagged bonds add entries that are not worth keeping. - """ - - def __init__(self) -> None: - self._compiled: dict[ - tuple[str, tuple[tuple[int, ...], ...]], ContractExpression - ] = {} - - def __call__(self, eq: str, *operands: NDArray) -> NDArray: - shapes = tuple(op.shape for op in operands) - expr = self._compiled.get((eq, shapes)) - if expr is None: - expr = self._compiled[eq, shapes] = contract_expression(eq, *shapes) - return expr(*operands) - - -class _Equations(NamedTuple): - """The einsum equations of one sweep, for a fixed stack depth.""" - - ltr: str - rtl_m: str - rtl_s: str - first: str - - -@cache -def _equations(depth: int) -> _Equations: - """Generate the sweep equations for a stack of ``depth`` layers. - - Layer ``i`` at a site carries ``(a_i, b_i, x_i, x_{i+1})``: left and right - bonds, then its upper and lower physical legs, so that consecutive layers share - ``x``. The output legs are ``x_0`` (up) and ``x_depth`` (down). - """ - symbols = map(get_symbol, count()) - sketch, eta_right, eta_left = next(symbols), next(symbols), next(symbols) - left = "".join(next(symbols) for _ in range(depth)) - right = "".join(next(symbols) for _ in range(depth)) - phys = [next(symbols) for _ in range(depth + 1)] - up, down = phys[0], phys[-1] - layers = ",".join( - f"{left[i]}{right[i]}{phys[i]}{phys[i + 1]}" for i in range(depth) - ) - return _Equations( - ltr=f"{sketch}{left},{sketch}{up}{down},{layers}->{sketch}{right}", - rtl_m=f"{sketch}{left},{layers},{eta_right}{right}->{eta_right}{up}{down}{sketch}", - rtl_s=f"{eta_left}{eta_right}{up}{down},{layers},{eta_right}{right}->{eta_left}{left}", - first=f"{layers},{eta_right}{right}->{left}{eta_right}{up}{down}", - ) +def _ranges(n: int, batch: int) -> Iterator[tuple[int, int]]: + """Split ``range(n)`` into consecutive ``(lo, hi)`` batches.""" + for lo in range(0, n, batch): + yield lo, min(lo + batch, n) def sweep( - layers: Sequence[Sequence[NDArray]], + layers: Sequence[Sequence[Site]], kind: TrainKind, chi_out: int, prng: np.random.Generator, @@ -101,12 +64,15 @@ def sweep( *, cutoff: float = 0.0, dtype: DTypeLike, + resources: Resources | None = None, ) -> list[NDArray]: """Contract and compress a stack in ket form with one SRC sweep. Args: layers: MPOs, optionally followed by one MPS, all with the same number - (at least three) of sites and matching physical legs. + (at least three) of sites and matching physical legs. Sites may be any + array-likes with ``shape``, ``dtype``, ``ndim`` and ``np.asarray`` + support; each is read only when the sweep reaches it. kind: The kind of the contracted train. 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 @@ -114,47 +80,141 @@ def sweep( xp: Array module (``numpy`` or ``cupy``). cutoff: Relative singular-value cutoff for adaptive bond truncation. dtype: The data type of the sketches. + resources: Memory budgets and scratch space; detected when ``None``. Returns: The site arrays of the compressed train in right-canonical form, as numpy arrays. """ - depth = len(layers) - n_sites = len(layers[0]) - eqs = _equations(depth) - contract = _Contractions() - # sites[j] holds the padded tensors of every layer at site j. - sites = list(zip(*(pad([xp.asarray(a) for a in layer]) for layer in layers))) - # Guarded: the max walks every site, wasted work unless it is logged. + shapes = padded_shapes(layers) + # 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)) + budgets = resolve_budgets(resources, xp) + plan = make_plan(shapes, site_bytes(layers), chi_out, work, budgets) + # Guarded: the per-site lists walk the whole plan, wasted work unless logged. if logger.isEnabledFor(logging.DEBUG): logger.debug( - "Largest environment: %d elements", - chi_out - * max(prod(t.shape[1] for t in sites[j]) for j in range(n_sites - 1)), + "SRC plan: prefetch=%s, device peak=%d B, host peak=%d B, disk=%d B, " + "scratch=%s, tiers=%s, batches (env, sketch, project)=%s", + plan.prefetch, + plan.device_peak, + plan.host_peak, + plan.disk_bytes, + budgets.scratch_dir if plan.disk_bytes else None, + [site.tier for site in plan.sites], + [ + (site.env_batch, site.sketch_batch, site.project_batch) + for site in plan.sites + ], + ) + env_shapes = [(chi_out, *(s[0] for s in site)) for site in shapes] + kernels = SiteKernels(len(layers)) + with ( + device_pool_limit(xp, budgets.device_cap), + SiteSource(layers, xp, depth=plan.prefetch) as source, + EnvironmentStore( + plan, env_shapes, work, xp, budgets.scratch_dir, copy_stream=new_stream(xp) + ) as store, + ): + tms = perf_counter_ns() + _left_to_right( + kernels, source, store, plan, chi_out=chi_out, prng=prng, xp=xp, dtype=dtype + ) + logger.debug("Left-to-right sweep: %.3f s", (perf_counter_ns() - tms) * 1e-9) + tms = perf_counter_ns() + eta = _right_to_left( + kernels, + source, + store, + plan, + chi_out=chi_out, + cutoff=cutoff, + xp=xp, + dtype=work, + ) + logger.debug("Right-to-left sweep: %.3f s", (perf_counter_ns() - tms) * 1e-9) + logger.debug( + "SRC stalls: sites %.3f s, environments %.3f s", + source.stall_seconds, + store.stall_seconds, ) + logger.debug("Device pool: %d B", device_pool_bytes(xp)) + return [to_numpy(site) for site in unpad(eta, kind)] + - tms = perf_counter_ns() - # C[j] is the sketched environment of sites 0 .. j-1. - C = [xp.ones((chi_out,) + (1,) * depth, dtype=dtype)] +def _left_to_right( + kernels: SiteKernels, + source: SiteSource, + store: EnvironmentStore, + plan: Plan, + *, + chi_out: int, + prng: np.random.Generator, + xp: ModuleType, + dtype: DTypeLike, +) -> None: + """Build the environments ``C_1 .. C_{n-1}`` into the store.""" + depth = len(kernels.eqs.ltr.split(",")) - 2 + first_env = xp.ones((chi_out,) + (1,) * depth, dtype=dtype) + n_sites = len(source) for j in range(n_sites - 1): - up, down = sites[j][0].shape[2], sites[j][-1].shape[3] + cores = source[j] + if plan.prefetch: + source.prefetch(j + 1) + up, down = cores[0].shape[2], cores[-1].shape[3] omega = gaussian_sketch(prng, (chi_out, up, down), dtype, xp) - C.append(contract(eqs.ltr, C[j], omega, *sites[j])) - logger.debug("Left-to-right sweep: %.3f s", (perf_counter_ns() - tms) * 1e-9) + batches = list(_ranges(chi_out, plan.sites[j].env_batch)) + if j > 0: + store.prefetch(j, batches) + for lo, hi in batches: + env = first_env[lo:hi] if j == 0 else store.get(j, lo, hi) + store.put(j + 1, lo, hi, kernels.env(env, omega[lo:hi], cores)) + + +def _right_to_left( + kernels: SiteKernels, + source: SiteSource, + store: EnvironmentStore, + plan: Plan, + *, + chi_out: int, + cutoff: float, + xp: ModuleType, + dtype: np.dtype, +) -> list[NDArray]: + """Build the output cores, host-side, from the last site to the first. - tms = perf_counter_ns() + ``dtype`` is the working dtype of the sweep, that of the environments. + """ + depth = len(kernels.eqs.ltr.split(",")) - 2 + n_sites = len(source) eta_reversed: list[NDArray] = [] - S = xp.ones((1,) * (depth + 1), dtype=dtype) + proj = xp.ones((1,) * (depth + 1), dtype=dtype) for j in range(n_sites - 1, 0, -1): - # C[-1] is C[j] here; popping it frees each environment once used. - M = contract(eqs.rtl_m, C.pop(), *sites[j], S) - rows = M.shape[0] * M.shape[1] * M.shape[2] - Q = truncated_qr(M.reshape(rows, chi_out), cutoff, xp) - eta_j = Q.reshape(*M.shape[:3], Q.shape[1]).transpose(3, 0, 1, 2) - S = contract(eqs.rtl_s, eta_j.conj(), *sites[j], S) - eta_reversed.append(eta_j) - first = contract(eqs.first, *sites[0], S) - eta = [first.reshape(1, *first.shape[depth:]), *reversed(eta_reversed)] - logger.debug("Right-to-left sweep: %.3f s", (perf_counter_ns() - tms) * 1e-9) - - return [to_numpy(site) for site in unpad(eta, kind)] + cores = source[j] + if plan.prefetch: + source.prefetch(j - 1) + site = plan.sites[j] + batches = list(_ranges(chi_out, site.sketch_batch)) + store.prefetch(j, batches) + up, down = cores[0].shape[2], cores[-1].shape[3] + sketch = xp.empty((proj.shape[0], up, down, chi_out), dtype=dtype) + for lo, hi in batches: + sketch[..., lo:hi] = kernels.sketch(store.get(j, lo, hi), cores, proj) + store.drop(j) + rows = sketch.shape[0] * up * down + Q = truncated_qr(sketch.reshape(rows, chi_out), cutoff, xp) + del sketch + eta_j = Q.reshape(proj.shape[0], up, down, Q.shape[1]).transpose(3, 0, 1, 2) + new_proj = xp.empty((Q.shape[1], *(c.shape[0] for c in cores)), dtype=dtype) + for lo, hi in _ranges(Q.shape[1], site.project_batch): + new_proj[lo:hi] = kernels.project(eta_j[lo:hi].conj(), cores, proj) + proj = new_proj + eta_reversed.append(to_numpy(eta_j)) + cores = source[0] + 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) + 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 eb0b36a..5e69020 100644 --- a/src/src_method/_tensor_train.py +++ b/src/src_method/_tensor_train.py @@ -16,7 +16,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Literal +from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeAlias import numpy as np from opt_einsum import contract @@ -41,21 +41,58 @@ TrainKind = Literal["mps", "mpo"] + +class SiteLike(Protocol): + """A site tensor that may be read lazily, such as a zarr or HDF5 dataset. + + Its shape, dtype and rank are known at once; the data is read only by + ``np.asarray``. NumPy and CuPy arrays and ``np.memmap`` all qualify. + """ + + @property + def shape(self) -> tuple[int, ...]: + """The shape of the site.""" + ... + + @property + def dtype(self) -> np.dtype: + """The dtype of the site.""" + ... + + @property + def ndim(self) -> int: + """The rank of the site, known without reading it.""" + ... + + def __array__(self, dtype: Any = None, copy: bool | None = None) -> np.ndarray: # noqa: ANN401, FBT001 (the NumPy array protocol) + """Read the site into a NumPy array.""" + ... + + +# A site as the public entry points accept it: an array or a lazily read array-like. +Site: TypeAlias = "NDArray | SiteLike" + __all__ = [ "MIN_SRC_SITES", + "Site", + "SiteLike", + "SwappedLegs", "TrainKind", "check_exact_supported", "exact_compress", "exact_stack", "infer_kind", + "known_kind", "normalize_stack", "pad", + "pad_site", + "padded_shape", "transpose_mpo", "unpad", ] -def infer_kind(arrays: Sequence[NDArray]) -> TrainKind | None: +def infer_kind(arrays: Sequence[Site]) -> TrainKind | None: """Classify a tensor train from the rank of its first site tensor. A boundary site carries one bond index plus either a single physical @@ -77,6 +114,25 @@ def infer_kind(arrays: Sequence[NDArray]) -> TrainKind | None: return None +def known_kind(arrays: Sequence[Site]) -> TrainKind: + """Classify a train already validated by `normalize_stack`; see `infer_kind`. + + Args: + arrays: The site tensors of the train. + + Returns: + ``"mps"`` or ``"mpo"``. + + Raises: + TypeError: If the layout is unrecognised. + """ + kind = infer_kind(arrays) + if kind is None: + msg = "Unrecognised train layout: expected an MPS or an MPO." + raise TypeError(msg) + return kind + + def check_exact_supported(n_sites: int) -> None: """Reject sub-``MIN_SRC_SITES`` trains the exact path cannot handle. @@ -138,6 +194,49 @@ def _truncated_svd(theta: NDArray, chi_out: int) -> tuple[NDArray, NDArray]: return U[:, :rank] * S[:rank], Vh[:rank] +def padded_shape( + shape: tuple[int, ...], kind: TrainKind, i: int, last: int +) -> tuple[int, int, int, int]: + """Return the bulk ``(l, r, u, d)`` shape that `pad_site` gives a site. + + Args: + shape: The unpadded shape of site ``i``. + kind: The kind of the train the site belongs to. + i: The position of the site. + last: The position of the last site of the train. + + Returns: + The padded shape. + """ + dims = [*shape, 1] if kind == "mps" else list(shape) + if i == 0: + dims.insert(0, 1) + if i == last: + dims.insert(1, 1) + left, right, up, down = dims + return left, right, up, down + + +def pad_site(site: NDArray, kind: TrainKind, i: int, last: int) -> NDArray: + """View one site as a bulk MPO tensor ``(l, r, u, d)``; see `pad`. + + Args: + site: The site tensor at position ``i``. + kind: The kind of the train the site belongs to. + i: The position of the site. + last: The position of the last site of the train. + + Returns: + The rank-4 view. + """ + view = site[..., None] if kind == "mps" else site + if i == 0: + view = view[None] + if i == last: + view = view[:, None] + return view + + def pad(train: Sequence[NDArray]) -> list[NDArray]: """View every site of a train as a bulk MPO tensor ``(l, r, u, d)``. @@ -151,17 +250,9 @@ def pad(train: Sequence[NDArray]) -> list[NDArray]: Returns: The rank-4 views, one per site. """ - kind = infer_kind(train) + kind = known_kind(train) last = len(train) - 1 - padded = [] - for i, site in enumerate(train): - view = site[..., None] if kind == "mps" else site - if i == 0: - view = view[None] - if i == last: - view = view[:, None] - padded.append(view) - return padded + return [pad_site(site, kind, i, last) for i, site in enumerate(train)] def unpad(train: Sequence[NDArray], kind: TrainKind) -> list[NDArray]: @@ -186,14 +277,48 @@ def unpad(train: Sequence[NDArray], kind: TrainKind) -> list[NDArray]: return unpadded -def transpose_mpo(train: Sequence[NDArray]) -> list[NDArray]: - """Transpose an MPO by swapping its ``u`` and ``d`` legs on every site (views).""" - return [site.swapaxes(-2, -1) for site in train] +class SwappedLegs: + """A site read lazily, with its ``u`` and ``d`` legs swapped on reading. + + Stands in for ``site.swapaxes(-2, -1)`` when ``site`` is a lazily loaded + array-like (a zarr or HDF5 dataset) that has no ``swapaxes``: the shape is known + at once and the data is only read by ``np.asarray``. + """ + + def __init__(self, site: Site) -> None: + """Wrap a lazily loaded site. + + Args: + site: The site to transpose. + """ + self._site = site + shape = tuple(site.shape) + self.shape = (*shape[:-2], shape[-1], shape[-2]) + self.dtype = np.dtype(site.dtype) + self.ndim = len(shape) + + def __array__(self, dtype: Any = None, copy: bool | None = None) -> np.ndarray: # noqa: ANN401, FBT001 + """Read the site and swap its physical legs.""" + del copy + swapped = np.asarray(self._site).swapaxes(-2, -1) + return swapped if dtype is None else swapped.astype(dtype) + + +def transpose_mpo(train: Sequence[Site]) -> list[Site]: + """Transpose an MPO by swapping its ``u`` and ``d`` legs on every site. + + Arrays give views; lazily loaded sites without ``swapaxes`` give `SwappedLegs`. + """ + transposed: list[Site] = [] + for site in train: + swapaxes = getattr(site, "swapaxes", None) + transposed.append(swapaxes(-2, -1) if callable(swapaxes) else SwappedLegs(site)) + return transposed def normalize_stack( - trains: Sequence[Sequence[NDArray]], -) -> tuple[list[Sequence[NDArray]], TrainKind]: + trains: Sequence[Sequence[Site]], +) -> tuple[list[Sequence[Site]], TrainKind]: """Validate a stack and rewrite it in ket form. A stack ``T_1 . T_2 . ... . T_m`` is contracted along the physical legs, the @@ -256,7 +381,7 @@ def _check_roles(kinds: Sequence[TrainKind | None]) -> None: def _check_physical_dims( - trains: Sequence[Sequence[NDArray]], kinds: Sequence[TrainKind | None] + trains: Sequence[Sequence[Site]], kinds: Sequence[TrainKind | None] ) -> None: """Check that the legs joined between adjacent trains agree at every site.""" for i in range(len(trains) - 1): @@ -272,7 +397,7 @@ def _check_physical_dims( def exact_stack( - layers: Sequence[Sequence[NDArray]], chi_out: int, kind: TrainKind + layers: Sequence[Sequence[Site]], chi_out: int, kind: TrainKind ) -> list[NDArray]: """Contract and compress a two-site stack exactly. diff --git a/src/src_method/_validation.py b/src/src_method/_validation.py index ed8cbbb..bfcb566 100644 --- a/src/src_method/_validation.py +++ b/src/src_method/_validation.py @@ -14,7 +14,7 @@ if TYPE_CHECKING: from collections.abc import Sequence - from numpy.typing import NDArray + from ._tensor_train import Site __all__ = ["validate_chi_out", "validate_cutoff", "validate_open_boundary"] @@ -61,7 +61,7 @@ def validate_cutoff(cutoff: float) -> None: def validate_open_boundary( - train: Sequence[NDArray], boundary_ndim: int, index: int + train: Sequence[Site], boundary_ndim: int, index: int ) -> None: """Check that a train has the layout of an open-boundary MPS or MPO. diff --git a/src/src_method/apply.py b/src/src_method/apply.py index 7ea086e..1fd6969 100644 --- a/src/src_method/apply.py +++ b/src/src_method/apply.py @@ -19,16 +19,20 @@ import numpy as np from numpy.typing import DTypeLike, NDArray + from ._plan import Resources + from ._tensor_train import Site + def apply( - left_tensor: Sequence[NDArray], - right_tensor: Sequence[NDArray], + left_tensor: Sequence[Site], + right_tensor: Sequence[Site], chi_out: int | np.integer, *, cutoff: float = 0.0, dtype: DTypeLike | None = None, seed: int | None = None, device: str = "cpu", + resources: Resources | None = None, ) -> list[NDArray]: """Applies the Successive Randomized Compression (SRC) algorithm. @@ -55,6 +59,7 @@ def apply( seed: An optional seed for the random number generator. device: ``"cpu"`` (default, numpy) or ``"gpu"`` (cupy). Requires the optional ``cupy`` dependency for GPU execution. + resources: Memory budgets and scratch space; see `src_method.stack.src`. Returns: The site arrays of the compressed tensor network (MPS or MPO). @@ -86,4 +91,5 @@ def apply( dtype=dtype, seed=seed, device=device, + resources=resources, ) diff --git a/src/src_method/compress.py b/src/src_method/compress.py index ef37d84..46e7ffe 100644 --- a/src/src_method/compress.py +++ b/src/src_method/compress.py @@ -18,15 +18,19 @@ import numpy as np from numpy.typing import DTypeLike, NDArray + from ._plan import Resources + from ._tensor_train import Site + def compress( - tensor: Sequence[NDArray], + tensor: Sequence[Site], chi_out: int | np.integer, *, cutoff: float = 0.0, dtype: DTypeLike | None = None, seed: int | None = None, device: str = "cpu", + resources: Resources | None = None, ) -> list[NDArray]: """Applies the Successive Randomized Compression (SRC) algorithm. @@ -51,6 +55,7 @@ def compress( seed: An optional seed for the random number generator. device: ``"cpu"`` (default, numpy) or ``"gpu"`` (cupy). Requires the optional ``cupy`` dependency for GPU execution. + resources: Memory budgets and scratch space; see `src_method.stack.src`. Returns: The site arrays of the compressed tensor network (MPS or MPO). @@ -65,5 +70,11 @@ def compress( ImportError: If ``device="gpu"`` but cupy is not installed. """ return src( - tensor, chi_out=chi_out, cutoff=cutoff, dtype=dtype, seed=seed, device=device + tensor, + chi_out=chi_out, + cutoff=cutoff, + dtype=dtype, + seed=seed, + device=device, + resources=resources, ) diff --git a/src/src_method/stack.py b/src/src_method/stack.py index 9b32baa..631e314 100644 --- a/src/src_method/stack.py +++ b/src/src_method/stack.py @@ -26,6 +26,9 @@ import numpy as np from numpy.typing import DTypeLike, NDArray + from ._plan import Resources + from ._tensor_train import Site + logger = logging.getLogger(__name__) LOG_WARN_SMALL = ( @@ -35,12 +38,13 @@ def src( - *trains: Sequence[NDArray], + *trains: Sequence[Site], chi_out: int | np.integer, cutoff: float = 0.0, dtype: DTypeLike | None = None, seed: int | None = None, device: str = "cpu", + resources: Resources | None = None, ) -> list[NDArray]: """Contract a stack of tensor trains and compress the result with SRC. @@ -66,7 +70,10 @@ def src( two or three Trotter layers). Apply anything else pairwise. Args: - *trains: The site arrays of each train, in mathematical order. + *trains: The site arrays of each train, in mathematical order. A site may + also be any array-like with ``shape``, ``dtype``, ``ndim`` and + ``np.asarray`` support (`SiteLike`), such as ``np.memmap`` or a zarr or + HDF5 dataset; it is then read only when the sweep reaches it. chi_out: The desired maximum bond dimension of the output train. cutoff: Relative singular-value cutoff for adaptive bond truncation. When positive, bonds are trimmed to their effective rank by @@ -81,6 +88,9 @@ def src( seed: An optional seed for the random number generator. device: ``"cpu"`` (default, numpy) or ``"gpu"`` (cupy). Requires the optional ``cupy`` dependency for GPU execution. + resources: Memory budgets and scratch space for the sweep (see + `Resources`); every budget left unset is detected. Ignored for + two-site stacks. Returns: The site arrays of the compressed train (MPS or MPO), in right-canonical @@ -123,6 +133,7 @@ def src( xp, cutoff=cutoff, dtype=sketch_dtype(dtype, *layers), + resources=resources, ) logger.debug("SRC complete") return result diff --git a/src/src_method/utils/__init__.py b/src/src_method/utils/__init__.py index b78556d..c3e661c 100644 --- a/src/src_method/utils/__init__.py +++ b/src/src_method/utils/__init__.py @@ -2,14 +2,44 @@ from __future__ import annotations -from ._backend import default_rng, gaussian_sketch, get_xp, sketch_dtype, to_numpy +from ._backend import ( + NullEvent, + NullStream, + current_stream, + default_rng, + device_memory, + device_pool_bytes, + device_pool_limit, + gaussian_sketch, + get_xp, + host_memory_available, + is_host, + new_stream, + pinned_empty, + sketch_dtype, + to_device_async, + to_host_async, + to_numpy, +) from .linalg import truncated_qr __all__ = [ + "NullEvent", + "NullStream", + "current_stream", "default_rng", + "device_memory", + "device_pool_bytes", + "device_pool_limit", "gaussian_sketch", "get_xp", + "host_memory_available", + "is_host", + "new_stream", + "pinned_empty", "sketch_dtype", + "to_device_async", + "to_host_async", "to_numpy", "truncated_qr", ] diff --git a/src/src_method/utils/_backend.py b/src/src_method/utils/_backend.py index a0f49ec..d249412 100644 --- a/src/src_method/utils/_backend.py +++ b/src/src_method/utils/_backend.py @@ -1,23 +1,30 @@ """Array-module backend selection for CPU (numpy) and GPU (cupy). -Kept intentionally minimal: a single resolver returns the appropriate -array module, a PRNG factory, and a host-transfer helper. All hot-loop -code paths receive an ``xp`` module and call ``xp.linalg.*`` / -``xp.asarray`` directly, so backend selection adds zero per-op overhead. +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. """ from __future__ import annotations -from typing import TYPE_CHECKING +import os +from contextlib import contextmanager +from pathlib import Path +from typing import TYPE_CHECKING, Any import numpy as np if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import Iterator, Sequence from types import ModuleType from numpy.typing import DTypeLike, NDArray + from src_method._tensor_train import Site, SiteLike + def get_xp(device: str) -> ModuleType: """Return the array module for the requested device. @@ -54,8 +61,11 @@ def default_rng( return np.random.default_rng(seed) -def to_numpy(arr: NDArray) -> np.ndarray: - """Bring an array onto the host as a numpy array (no-op for numpy).""" +def to_numpy(arr: NDArray | SiteLike) -> np.ndarray: + """Bring an array, or a lazily read site, onto the host as a numpy array. + + A no-op for numpy arrays. + """ if isinstance(arr, np.ndarray): return arr # cupy.ndarray exposes .get(); fall back to np.asarray for other dispatchers. @@ -63,7 +73,176 @@ def to_numpy(arr: NDArray) -> np.ndarray: return get() if callable(get) else np.asarray(arr) -def sketch_dtype(dtype: DTypeLike | None, *inputs: Sequence[NDArray]) -> np.dtype: +class NullEvent: + """Stand-in for a CUDA event on the host backend: already complete.""" + + def synchronize(self) -> None: + """Return at once: host work is synchronous.""" + + +class NullStream: + """Stand-in for a CUDA stream on the host backend: work runs synchronously.""" + + def record(self) -> NullEvent: + """Return an event that is already complete.""" + return NullEvent() + + def wait_event(self, event: NullEvent) -> None: + """Return at once: there is nothing to wait for.""" + + def synchronize(self) -> None: + """Return at once: host work is synchronous.""" + + +def is_host(xp: ModuleType) -> bool: + """Whether ``xp`` is the host backend (numpy).""" + return xp is np + + +def new_stream(xp: ModuleType) -> Any: # noqa: ANN401 (a cupy or null stream) + """Return a non-blocking stream for transfers, or a `NullStream` on the host.""" + if is_host(xp): + return NullStream() + return xp.cuda.Stream(non_blocking=True) + + +def current_stream(xp: ModuleType) -> Any: # noqa: ANN401 (a cupy or null stream) + """Return the stream kernels run on, or a `NullStream` on the host.""" + if is_host(xp): + return NullStream() + return xp.cuda.get_current_stream() + + +def pinned_empty(n_bytes: int, xp: ModuleType) -> np.ndarray: + """Allocate a flat byte buffer, page-locked on the GPU backend. + + Args: + n_bytes: The size of the buffer. + xp: Array module (``numpy`` or ``cupy``). + + Returns: + A ``uint8`` host array of ``n_bytes`` elements. + """ + if is_host(xp): + return np.empty(n_bytes, dtype=np.uint8) + import cupyx # noqa: PLC0415 (lazy: optional dependency) + + return cupyx.empty_pinned(n_bytes, dtype=np.uint8) + + +def to_device_async(host: np.ndarray, xp: ModuleType, stream: Any) -> NDArray: # 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 + buffer that is about to be reused. + + Args: + host: The source; page-locked for the copy to be asynchronous. + xp: Array module (``numpy`` or ``cupy``). + stream: The stream that performs the copy. + + Returns: + The new device array. + """ + if is_host(xp): + return host.copy() + device = xp.empty(host.shape, dtype=host.dtype) + device.set(host, stream=stream) + return device + + +def to_host_async(device: NDArray, 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. + + Args: + device: The source. + out: The destination, of the same shape and dtype; page-locked for the + copy to be asynchronous. + stream: The stream that performs the copy. + """ + if isinstance(device, np.ndarray): + np.copyto(out, device) + return + device.get(stream=stream, out=out, blocking=False) + + +def device_memory(xp: ModuleType) -> tuple[int, int]: + """Return the device bytes available to the sweep and the device total. + + Available bytes are the free device memory plus the bytes cached, but unused, + by cupy's default memory pool. + + Args: + xp: The cupy module. + + Returns: + ``(available, total)`` in bytes. + """ + free, total = xp.cuda.runtime.memGetInfo() + return free + xp.get_default_memory_pool().free_bytes(), total + + +def device_pool_bytes(xp: ModuleType) -> int: + """Return the bytes held by cupy's default memory pool, or 0 on the host. + + The pool keeps freed blocks for reuse, so after a sweep this is its high-water + mark. + """ + if is_host(xp): + return 0 + return xp.get_default_memory_pool().total_bytes() + + +@contextmanager +def device_pool_limit(xp: ModuleType, budget: int | None) -> Iterator[None]: + """Cap cupy's default memory pool at ``budget`` bytes beyond its current use. + + An allocation past the cap fails at once instead of when another allocation + runs out. The previous limit is restored on exit. No-op on the host backend or + without a budget. + + Args: + xp: Array module (``numpy`` or ``cupy``). + budget: The bytes the pool may allocate on top of those in use, or + ``None`` for no cap. + + Yields: + Nothing. + """ + if is_host(xp) or budget is None: + yield + return + pool = xp.get_default_memory_pool() + previous = pool.get_limit() + pool.set_limit(size=pool.used_bytes() + budget) + try: + yield + finally: + pool.set_limit(size=previous) + + +def host_memory_available(meminfo: str = "/proc/meminfo") -> int: + """Return the host memory available for new allocations, in bytes. + + Reads ``MemAvailable`` on Linux and falls back to the free physical pages. + + Args: + meminfo: The path of the ``meminfo`` file. + + Returns: + The available bytes. + """ + path = Path(meminfo) + if path.exists(): + for line in path.read_text().splitlines(): + if line.startswith("MemAvailable:"): + return int(line.split()[1]) * 1024 + return os.sysconf("SC_AVPHYS_PAGES") * os.sysconf("SC_PAGE_SIZE") + + +def sketch_dtype(dtype: DTypeLike | None, *inputs: Sequence[Site]) -> np.dtype: """Resolve the sketch dtype. Args: diff --git a/tests/test_backend.py b/tests/test_backend.py new file mode 100644 index 0000000..d69acd3 --- /dev/null +++ b/tests/test_backend.py @@ -0,0 +1,74 @@ +"""Test the host-side behaviour of the backend helpers.""" + +import numpy as np + +from src_method.utils import ( + NullEvent, + NullStream, + current_stream, + device_pool_bytes, + device_pool_limit, + host_memory_available, + is_host, + new_stream, + pinned_empty, + to_device_async, + to_host_async, +) + + +def test_host_streams_are_null(): + stream = new_stream(np) + + assert is_host(np) + assert isinstance(stream, NullStream) + assert isinstance(current_stream(np), NullStream) + event = stream.record() + assert isinstance(event, NullEvent) + stream.wait_event(event) + event.synchronize() + stream.synchronize() + + +def test_pinned_empty_on_host_is_a_byte_buffer(): + buffer = pinned_empty(24, np) + + assert buffer.dtype == np.uint8 + assert buffer.shape == (24,) + + +def test_to_device_async_on_host_copies(): + host = np.arange(6.0).reshape(2, 3) + + device = to_device_async(host, np, NullStream()) + host[:] = 0 + + np.testing.assert_array_equal(device, np.arange(6.0).reshape(2, 3)) + + +def test_to_host_async_on_host_fills_out(): + out = np.empty((2, 3)) + + to_host_async(np.arange(6.0).reshape(2, 3).T.T, out, NullStream()) + + np.testing.assert_array_equal(out, np.arange(6.0).reshape(2, 3)) + + +def test_device_pool_limit_is_a_no_op_on_host(): + with device_pool_limit(np, 10): + pass + + +def test_device_pool_is_empty_on_host(): + assert device_pool_bytes(np) == 0 + + +def test_host_memory_reads_mem_available(tmp_path): + meminfo = tmp_path / "meminfo" + meminfo.write_text("MemTotal: 100 kB\nMemAvailable: 2048 kB\n") + + assert host_memory_available(str(meminfo)) == 2048 * 1024 + + +def test_host_memory_falls_back_without_meminfo(tmp_path): + assert host_memory_available(str(tmp_path / "missing")) > 0 diff --git a/tests/test_gpu_backend.py b/tests/test_gpu_backend.py index 420604a..6b4fdb6 100644 --- a/tests/test_gpu_backend.py +++ b/tests/test_gpu_backend.py @@ -5,11 +5,23 @@ from __future__ import annotations +from typing import Any + import numpy as np import pytest import quimb.tensor as qtn -from src_method import apply, compress, src +import src_method._sweep as sweep_module +from src_method import Resources, apply, compress, src +from src_method._kernels import SiteKernels +from src_method._plan import ( + GPU_MARGIN_FRACTION, + GPU_MARGIN_MIN, + Budgets, + Plan, + make_plan, +) +from src_method._sites import padded_shapes, site_bytes cupy = pytest.importorskip("cupy") @@ -181,3 +193,127 @@ def test_src_stack_gpu_matches_reference(device: str) -> None: ref = H1.apply(H2.apply(psi, compress=False), compress=False) np.testing.assert_allclose(ref.distance(out), 0.0, atol=1e-6) + + +# --------------------------------------- +# --- Budgets, batching and spilling --- +# --------------------------------------- + + +def random_mpo_arrays(bonds: list[int], rng: np.random.Generator) -> list[np.ndarray]: + """Complex Gaussian MPO with the given bonds and physical legs of 2.""" + + def site(*shape: int) -> np.ndarray: + return rng.normal(size=shape) + 1j * rng.normal(size=shape) + + lefts, rights = [None, *bonds], [*bonds, None] + return [ + site(*(b for b in (lb, rb) if b is not None), 2, 2) + for lb, rb in zip(lefts, rights) + ] + + +def dense_mpo(train: list[np.ndarray]) -> np.ndarray: + """Contract an MPO into a ``(U, D)`` matrix.""" + T = train[0] + for W in train[1:-1]: + T = np.einsum("aUD,abud->bUuDd", T, W) + T = T.reshape(W.shape[1], T.shape[1] * T.shape[2], T.shape[3] * T.shape[4]) + T = np.einsum("aUD,aud->UuDd", T, train[-1]) + return T.reshape(T.shape[0] * T.shape[1], T.shape[2] * T.shape[3]) + + +@pytest.fixture +def plans(monkeypatch: pytest.MonkeyPatch) -> list[Plan]: + """Record the plan of every sweep.""" + recorded: list[Plan] = [] + + def spy(*args: Any, **kwargs: Any) -> Plan: + recorded.append(make_plan(*args, **kwargs)) + return recorded[-1] + + monkeypatch.setattr(sweep_module, "make_plan", spy) + return recorded + + +def test_gpu_tiers_and_batches_match_cpu(tmp_path, plans: list[Plan]) -> None: + """Every tier and small batches on the GPU give the operator of the CPU run.""" + rng = np.random.default_rng(7) + stack = [random_mpo_arrays([4, 8, 8, 8, 4], rng) for _ in range(4)] + chi = 64 + shapes, sizes = padded_shapes(stack), site_bytes(stack) + roomy = make_plan( + shapes, + sizes, + chi, + np.complex128, + Budgets(10**10, 10**10, 10**12, tmp_path, unified=False), + ) + env = chi * 8**4 * 16 # one bulk environment + tight = Resources( + gpu_memory=roomy.device_peak // 2, + host_memory=roomy.host_peak + 5 * env, + scratch_dir=tmp_path, + ) + + cpu = src(*stack, chi_out=chi, seed=3, dtype=np.complex128) + gpu = src( + *stack, chi_out=chi, seed=3, dtype=np.complex128, device="gpu", resources=tight + ) + + assert {site.tier for site in plans[-1].sites} == {"device", "host", "disk"} + assert min(site.sketch_batch for site in plans[-1].sites[1:]) < chi + reference = dense_mpo(cpu) + error = np.linalg.norm(dense_mpo(gpu) - reference) / np.linalg.norm(reference) + assert error < 1e-10 + assert list(tmp_path.iterdir()) == [] + + +def test_gpu_pool_limit_is_restored(monkeypatch: pytest.MonkeyPatch) -> None: + """The cap on cupy's pool is lifted after the sweep, also after an error.""" + rng = np.random.default_rng(8) + stack = [random_mpo_arrays([2, 3, 3, 2], rng) for _ in range(2)] + pool = cupy.get_default_memory_pool() + previous = pool.get_limit() + + src(*stack, chi_out=4, seed=0, device="gpu") + assert pool.get_limit() == previous + + def boom(*_args: object) -> None: + msg = "boom" + raise RuntimeError(msg) + + monkeypatch.setattr(SiteKernels, "sketch", boom) + with pytest.raises(RuntimeError, match="boom"): + src(*stack, chi_out=4, seed=0, device="gpu") + assert pool.get_limit() == previous + + +def test_gpu_pool_cap_leaves_room_for_fragmentation(tmp_path) -> None: + """The pool is capped at the GPU budget plus the margin, not at the budget.""" + rng = np.random.default_rng(9) + stack = [random_mpo_arrays([2, 3, 3, 2], rng) for _ in range(2)] + pool = cupy.get_default_memory_pool() + limits: list[int] = [] + + def spy(*_args: object) -> None: + limits.append(pool.get_limit()) + msg = "spy" + raise RuntimeError(msg) + + budget = 10**8 + used = pool.used_bytes() + with pytest.MonkeyPatch.context() as mp: + mp.setattr(SiteKernels, "env", spy) + with pytest.raises(RuntimeError, match="spy"): + src( + *stack, + chi_out=4, + seed=0, + device="gpu", + resources=Resources(gpu_memory=budget, scratch_dir=tmp_path), + ) + + _, total = cupy.cuda.runtime.memGetInfo() + margin = max(int(GPU_MARGIN_FRACTION * total), GPU_MARGIN_MIN) + assert limits == [used + budget + margin] diff --git a/tests/test_kernels.py b/tests/test_kernels.py new file mode 100644 index 0000000..a3eb6fc --- /dev/null +++ b/tests/test_kernels.py @@ -0,0 +1,82 @@ +"""Test the batched site kernels and the peak-memory walk.""" + +import numpy as np +import pytest + +from src_method._kernels import SiteKernels, equations, peak_elements + + +def cores(rng, dtype, *, left=(3, 2), right=(4, 3), up=2, mid=3, down=2): + """Two padded layers: (l, r, u, x) and (l, r, x, d).""" + + def draw(*shape): + out = rng.normal(size=shape) + if np.issubdtype(dtype, np.complexfloating): + out = out + 1j * rng.normal(size=shape) + return out.astype(dtype) + + return ( + draw(left[0], right[0], up, mid), + draw(left[1], right[1], mid, down), + ) + + +def batched(fn, n, batch, axis): + parts = [fn(lo, min(lo + batch, n)) for lo in range(0, n, batch)] + return np.concatenate(parts, axis=axis) + + +@pytest.mark.parametrize("dtype", [np.float64, np.complex128]) +@pytest.mark.parametrize("batch", [1, 3, 7]) +def test_batched_kernels_match_unbatched(dtype, batch): + rng = np.random.default_rng(0) + k = SiteKernels(2) + site = cores(rng, dtype) + chi, eta = 7, 5 + env = rng.normal(size=(chi, 3, 2)).astype(dtype) + omega = rng.normal(size=(chi, 2, 2)).astype(dtype) + proj = rng.normal(size=(eta, 4, 3)).astype(dtype) + out_core = rng.normal(size=(chi, eta, 2, 2)).astype(dtype) + + full_env = k.env(env, omega, site) + full_sketch = k.sketch(env, site, proj) + full_proj = k.project(out_core, site, proj) + full_first = k.first(site, proj) + + np.testing.assert_allclose( + batched(lambda lo, hi: k.env(env[lo:hi], omega[lo:hi], site), chi, batch, 0), + full_env, + rtol=1e-12, + ) + np.testing.assert_allclose( + batched(lambda lo, hi: k.sketch(env[lo:hi], site, proj), chi, batch, 3), + full_sketch, + rtol=1e-12, + ) + np.testing.assert_allclose( + batched(lambda lo, hi: k.project(out_core[lo:hi], site, proj), chi, batch, 0), + full_proj, + rtol=1e-12, + ) + np.testing.assert_allclose( + batched(lambda lo, hi: k.first(site, proj[lo:hi]), eta, batch, 2), + full_first, + rtol=1e-12, + ) + + +def test_equations_depth_one(): + eqs = equations(1) + + assert eqs.ltr == "ad,afg,defg->ae" + assert eqs.first == "defg,be->dbfg" + + +def test_peak_elements_of_a_matrix_chain(): + # (2x3)(3x4)(4x5): the path contracts the first pair into a 2x4 intermediate, + # holding its two operands (6 + 12) and the output (8), then the last pair, + # holding that intermediate twice (as live and as operand), the 4x5 operand + # and the 2x5 output: 8 + 8 + 20 + 10 = 46. + shapes = ((2, 3), (3, 4), (4, 5)) + + assert peak_elements("ab,bc,cd->ad", shapes) == 46 diff --git a/tests/test_plan.py b/tests/test_plan.py new file mode 100644 index 0000000..67725c2 --- /dev/null +++ b/tests/test_plan.py @@ -0,0 +1,253 @@ +"""Test the memory planner: sizes, budgets, batches and tiers.""" + +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._plan import ( + GEMM_MULTIPLE, + Budgets, + make_plan, + parse_size, + resolve_budgets, +) + +GB = 10**9 + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + (123, 123), + ("36GB", 36 * 10**9), + ("36GiB", 36 * 2**30), + ("1.5 kB", 1500), + ("512", 512), + ("2mib", 2 * 2**20), + ], +) +def test_parse_size(value, expected): + assert parse_size(value) == expected + + +@pytest.mark.parametrize("value", [-1, "36 GBs", "lots", ""]) +def test_parse_size_rejects_values(value): + with pytest.raises(ValueError, match="Expected"): + parse_size(value) + + +@pytest.mark.parametrize("value", [True, 1.5, None]) +def test_parse_size_rejects_types(value): + with pytest.raises(TypeError, match="Expected"): + parse_size(value) + + +def test_resources_validate_at_construction(): + with pytest.raises(ValueError, match="size string"): + Resources(gpu_memory="plenty") + + +def test_resolve_budgets_explicit_on_host(tmp_path): + budgets = resolve_budgets(Resources(host_memory="2GB", scratch_dir=tmp_path), np) + + assert budgets.unified + assert budgets.host == budgets.device == 2 * GB + assert budgets.scratch_dir == tmp_path + assert budgets.disk > 0 + + +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) + + assert budgets.host == 9 * GB + + +def test_resolve_budgets_detects_device_memory(monkeypatch, tmp_path): + fake_xp = ModuleType("fake_xp") + monkeypatch.setattr(plan_module, "is_host", lambda _xp: False) + monkeypatch.setattr(plan_module, "device_memory", lambda _xp: (30 * GB, 40 * GB)) + + budgets = resolve_budgets( + Resources(host_memory="1GB", scratch_dir=tmp_path), fake_xp + ) + + assert not budgets.unified + assert budgets.device == 26 * GB # minus max(10% of 40 GB, 1 GiB) + assert budgets.device_cap == 30 * GB # the pool may grow into the margin + + +def test_resolve_budgets_caps_explicit_device_memory_with_margin(monkeypatch, tmp_path): + fake_xp = ModuleType("fake_xp") + monkeypatch.setattr(plan_module, "is_host", lambda _xp: False) + monkeypatch.setattr(plan_module, "device_memory", lambda _xp: (30 * GB, 40 * GB)) + + budgets = resolve_budgets( + Resources(gpu_memory="10GB", host_memory="1GB", scratch_dir=tmp_path), + fake_xp, + ) + + assert budgets.device == 10 * GB + assert budgets.device_cap == 14 * GB + + +def test_resolve_budgets_leaves_the_host_uncapped(tmp_path): + budgets = resolve_budgets(Resources(host_memory="1GB", scratch_dir=tmp_path), np) + + assert budgets.device_cap is None + + +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) + + assert budgets.disk == 95 * GB + + +def mpo_stack_shapes(n_sites, bonds, phys=4): + """Padded shapes of a stack of MPOs, one bond dimension per layer.""" + + def shape(j, bond): + return (1 if j == 0 else bond, 1 if j == n_sites - 1 else bond, phys, phys) + + return [tuple(shape(j, bond) for bond in bonds) for j in range(n_sites)] + + +def site_bytes(shapes, itemsize=16): + return [sum(int(np.prod(s)) * itemsize for s in site) for site in shapes] + + +def budgets(device, host=None, disk=10**15, *, unified=False): + return Budgets( + device, device if host is None else host, disk, Path("/scratch"), unified + ) + + +def test_small_problem_is_one_batch_on_the_device(): + shapes = mpo_stack_shapes(6, [2, 3, 2]) + + plan = make_plan(shapes, site_bytes(shapes), 64, np.complex128, budgets(GB)) + + assert plan.prefetch == 1 + assert all(site.tier == "device" for site in plan.sites) + assert plan.sites[0].env_batch == 64 + assert plan.sites[-1].env_batch == 0 + assert plan.sites[0].sketch_batch == 0 + assert all(site.sketch_batch == 64 for site in plan.sites[1:]) + assert all(site.project_batch == 64 for site in plan.sites) + assert plan.disk_bytes == 0 + + +def test_tight_budget_shrinks_batches_to_gemm_multiples(): + shapes = mpo_stack_shapes(6, [4, 4, 64, 4]) + chi = 512 + loose = make_plan(shapes, site_bytes(shapes), chi, np.complex128, budgets(10 * GB)) + + tight = make_plan( + shapes, site_bytes(shapes), chi, np.complex128, budgets(loose.device_peak // 4) + ) + + batches = [s.sketch_batch for s in tight.sites[1:]] + assert min(batches) < chi + assert all(b % GEMM_MULTIPLE == 0 for b in batches if b >= GEMM_MULTIPLE) + assert tight.device_peak <= loose.device_peak // 4 + + +def test_tiers_go_newest_first(): + shapes = mpo_stack_shapes(8, [4, 4, 64, 4]) + chi = 256 + env = chi * 4 * 4 * 64 * 4 * 16 # one bulk environment, complex128 + roomy = make_plan(shapes, site_bytes(shapes), chi, np.complex128, budgets(10 * GB)) + + seen = set() + for extra in range(8): + plan = make_plan( + shapes, + site_bytes(shapes), + chi, + np.complex128, + budgets(roomy.device_peak - extra * env, host=roomy.host_peak + 7 * env), + ) + tiers = [site.tier for site in plan.sites[1:]] + # Oldest sites on the slowest tier: disk, then host, then device. + assert tiers == sorted(tiers, key=["disk", "host", "device"].index) + assert plan.disk_bytes == env * tiers.count("disk") + seen.update(tiers) + assert seen == {"device", "host", "disk"} + + +def test_unified_memory_has_no_host_tier(): + shapes = mpo_stack_shapes(8, [4, 4, 64, 4]) + chi = 256 + roomy = make_plan( + shapes, site_bytes(shapes), chi, np.complex128, budgets(10 * GB, unified=True) + ) + + plan = make_plan( + shapes, + site_bytes(shapes), + chi, + np.complex128, + budgets(roomy.device_peak // 2, unified=True), + ) + + tiers = {site.tier for site in plan.sites[1:]} + assert "host" not in tiers + assert "disk" in tiers + + +def test_infeasible_site_names_the_site(): + shapes = mpo_stack_shapes(6, [4, 4, 64, 4]) + + with pytest.raises(MemoryError, match=r"Site \d+: the \w+ step needs \d+ bytes"): + make_plan(shapes, site_bytes(shapes), 256, np.complex128, budgets(10**6)) + + +def test_environments_must_fit_on_disk(): + shapes = mpo_stack_shapes(8, [4, 4, 64, 4]) + chi = 256 + roomy = make_plan(shapes, site_bytes(shapes), chi, np.complex128, budgets(10 * GB)) + work = roomy.device_peak - 7 * chi * 4 * 4 * 64 * 4 * 16 + + with pytest.raises(MemoryError, match="bytes on disk"): + make_plan( + shapes, + site_bytes(shapes), + chi, + np.complex128, + budgets(work, host=roomy.host_peak, disk=1000), + ) + + +def test_prefetch_is_dropped_before_giving_up(): + shapes = mpo_stack_shapes(6, [4, 4, 64, 4]) + chi = 64 + cores = max(site_bytes(shapes)) + with_prefetch = make_plan( + shapes, site_bytes(shapes), chi, np.complex128, budgets(10 * GB) + ) + # Remove about one site of cores from the smallest budget that fits one column. + minimum = with_prefetch.device_peak + lo, hi = 0, minimum + while lo < hi: # the smallest budget that plans with prefetching + mid = (lo + hi) // 2 + try: + make_plan(shapes, site_bytes(shapes), chi, np.complex128, budgets(mid)) + hi = mid + except MemoryError: + lo = mid + 1 + + plan = make_plan(shapes, site_bytes(shapes), chi, np.complex128, budgets(lo)) + assert plan.prefetch in {0, 1} + if plan.prefetch == 1: + dropped = make_plan( + shapes, site_bytes(shapes), chi, np.complex128, budgets(lo - cores // 2) + ) + assert dropped.prefetch == 0 diff --git a/tests/test_sites.py b/tests/test_sites.py new file mode 100644 index 0000000..f9a0364 --- /dev/null +++ b/tests/test_sites.py @@ -0,0 +1,122 @@ +"""Test the lazy site source.""" + +from collections.abc import Sequence + +import numpy as np +import pytest + +from src_method import src +from src_method._sites import SiteSource, padded_shapes, site_bytes +from src_method._tensor_train import pad + + +class CountingTrain(Sequence): + """A train that counts how often each site is read.""" + + def __init__(self, sites): + self.sites = sites + self.reads = [0] * len(sites) + + def __len__(self): + return len(self.sites) + + def __getitem__(self, j): + return CountingSite(self, j) + + +class CountingSite: + """A lazily read site: shape and dtype at once, data on ``np.asarray``.""" + + def __init__(self, train, j): + self.train, self.j = train, j + self.shape = train.sites[j].shape + self.dtype = train.sites[j].dtype + self.ndim = train.sites[j].ndim + + def __array__(self, dtype=None, copy=None): + self.train.reads[self.j] += 1 + return np.asarray(self.train.sites[self.j], dtype=dtype) + + +def random_mpo(n_sites, bond, rng, phys=2): + shapes = ( + [(bond, phys, phys)] + + [(bond, bond, phys, phys)] * (n_sites - 2) + + [(bond, phys, phys)] + ) + return [rng.normal(size=s) + 1j * rng.normal(size=s) for s in shapes] + + +def test_padded_shapes_and_bytes_without_reading(): + rng = np.random.default_rng(0) + train = CountingTrain(random_mpo(4, 3, rng)) + mps = [t[..., 0] for t in random_mpo(4, 2, rng)] + + shapes = padded_shapes([train, mps]) + + assert shapes == [ + tuple(t.shape for t in sites) for sites in zip(pad(train.sites), pad(mps)) + ] + assert site_bytes([train]) == [t.nbytes for t in train.sites] + assert train.reads == [0, 0, 0, 0] + + +@pytest.mark.parametrize("depth", [0, 1]) +def test_sites_are_read_when_requested(depth): + rng = np.random.default_rng(1) + train = CountingTrain(random_mpo(4, 3, rng)) + + with SiteSource([train], np, depth=depth) as source: + (core,) = source[2] + assert train.reads == [0, 0, 1, 0] + source.prefetch(3) + (last,) = source[3] + assert train.reads == [0, 0, 1, 1] + + np.testing.assert_array_equal(core, pad(train.sites)[2]) + np.testing.assert_array_equal(last, pad(train.sites)[3]) + + +def test_src_reads_each_site_once_per_pass(): + rng = np.random.default_rng(2) + train = CountingTrain(random_mpo(5, 3, rng)) + other = random_mpo(5, 2, rng) + + lazy = src(CountingTrain(other), train, chi_out=4, seed=0) + eager = src(other, train.sites, chi_out=4, seed=0) + + # Left-to-right reads sites 0..3, right-to-left sites 4..0. + assert train.reads == [2, 2, 2, 2, 1] + for a, b in zip(lazy, eager): + np.testing.assert_array_equal(a, b) + + +def test_src_on_memmaps(tmp_path): + rng = np.random.default_rng(3) + trains = [random_mpo(5, 3, rng), random_mpo(5, 2, rng)] + mapped = [] + for t, train in enumerate(trains): + sites = [] + for j, site in enumerate(train): + path = tmp_path / f"train{t}-site{j}.npy" + np.save(path, site) + sites.append(np.load(path, mmap_mode="r")) + mapped.append(sites) + + lazy = src(*mapped, chi_out=4, seed=0) + eager = src(*trains, chi_out=4, seed=0) + + for a, b in zip(lazy, eager): + np.testing.assert_array_equal(a, b) + + +def test_src_bra_stack_of_lazy_sites(): + rng = np.random.default_rng(4) + phi = [t[..., 0] for t in random_mpo(5, 2, rng)] + mpo = random_mpo(5, 3, rng) + + lazy = src(phi, CountingTrain(mpo), chi_out=4, seed=0) + eager = src(phi, mpo, chi_out=4, seed=0) + + for a, b in zip(lazy, eager): + np.testing.assert_array_equal(a, b) diff --git a/tests/test_stack.py b/tests/test_stack.py index 800b45d..693f7ca 100644 --- a/tests/test_stack.py +++ b/tests/test_stack.py @@ -5,7 +5,9 @@ import numpy as np import pytest -from src_method import apply, compress, src +import src_method._sweep as sweep_module +from src_method import Resources, apply, compress, src +from src_method._plan import make_plan from src_method._sweep import sweep # ------------- @@ -348,3 +350,58 @@ def test_benchmark_src_stack_depth3(benchmark): np.testing.assert_allclose( ref.distance(qtn.MatrixProductState(out)), 0.0, atol=1e-6 ) + + +# ---------------------------- +# --- Budgets and spilling --- +# ---------------------------- + + +@pytest.fixture +def plans(monkeypatch): + """Record the plan of every sweep.""" + recorded = [] + + def spy(*args, **kwargs): + recorded.append(make_plan(*args, **kwargs)) + return recorded[-1] + + monkeypatch.setattr(sweep_module, "make_plan", spy) + return recorded + + +def test_tiny_budget_batches_and_spills(rng, tmp_path, plans): + stack = [random_mpo([3, 4, 4, 4, 3], rng) for _ in range(4)] + tight = Resources(host_memory="1MB", scratch_dir=tmp_path) + + default = src(*stack, chi_out=16, seed=3, dtype=np.complex128) + spilled = src(*stack, chi_out=16, seed=3, dtype=np.complex128, resources=tight) + + assert all(site.tier == "device" for site in plans[0].sites) + assert "disk" in {site.tier for site in plans[1].sites} + assert min(site.sketch_batch for site in plans[1].sites[1:]) < 16 + # Rounding differs, and the output cores of an ill-conditioned sketch with it, + # but not the operator they represent. + assert rel_error(spilled, dense(default)) < 1e-10 + assert list(tmp_path.iterdir()) == [] + + +def test_budget_too_small_raises(rng): + stack = make_stack("A B psi", rng) + + with pytest.raises(MemoryError, match="working set exceeds the budget"): + src(*stack, chi_out=8, resources=Resources(host_memory="1kB")) + + +def test_apply_passes_resources(rng): + A, psi = make_stack("A psi", rng) + + with pytest.raises(MemoryError, match="working set exceeds the budget"): + apply(A, psi, chi_out=8, resources=Resources(host_memory="1kB")) + + +def test_compress_passes_resources(rng): + (A,) = make_stack("A", rng) + + with pytest.raises(MemoryError, match="working set exceeds the budget"): + compress(A, chi_out=8, resources=Resources(host_memory="1kB")) diff --git a/tests/test_store.py b/tests/test_store.py new file mode 100644 index 0000000..c9a22a7 --- /dev/null +++ b/tests/test_store.py @@ -0,0 +1,124 @@ +"""Test the environment store on every tier, on the host backend.""" + +import errno + +import numpy as np +import pytest + +import src_method._store as store_module +from src_method._plan import Plan, SitePlan +from src_method._store import EnvironmentStore +from src_method.utils import NullStream + +CHI = 10 +SHAPES = [(CHI, 1), (CHI, 3, 2), (CHI, 4, 2), (CHI, 2, 1)] + + +def make_plan(tier, *, env=4, sketch=3): + """Environments of sites 1-3 on ``tier``, written in batches of ``env``.""" + sites = [SitePlan(env, 0, 5, "device")] + sites += [SitePlan(env, sketch, 5, tier) for _ in range(2)] + sites += [SitePlan(0, sketch, 5, tier)] + return Plan(tuple(sites), 1, 0, 0, 0) + + +def ranges(n, batch): + return [(lo, min(lo + batch, n)) for lo in range(0, n, batch)] + + +def fill(store, rng, batch=4): + """Put random environments for sites 1-3 and return them.""" + envs = {} + for j in (1, 2, 3): + envs[j] = rng.normal(size=SHAPES[j]) + 1j * rng.normal(size=SHAPES[j]) + for lo, hi in ranges(CHI, batch): + store.put(j, lo, hi, envs[j][lo:hi]) + return envs + + +def open_store(tier, tmp_path, **batches): + return EnvironmentStore( + make_plan(tier, **batches), + SHAPES, + np.complex128, + np, + tmp_path, + copy_stream=NullStream(), + ) + + +@pytest.mark.parametrize("tier", ["device", "host", "disk"]) +def test_round_trip_in_other_batches(tier, tmp_path): + rng = np.random.default_rng(0) + with open_store(tier, tmp_path) as store: + envs = fill(store, rng) + for j in (3, 2, 1): + batches = ranges(CHI, 3) # read in batches of 3, written in batches of 4 + store.prefetch(j, batches) + got = np.concatenate([store.get(j, lo, hi) for lo, hi in batches]) + np.testing.assert_array_equal(got, envs[j]) + store.drop(j) + + +def test_get_without_prefetch(tmp_path): + rng = np.random.default_rng(1) + with open_store("disk", tmp_path) as store: + envs = fill(store, rng) + + np.testing.assert_array_equal(store.get(2, 7, 10), envs[2][7:10]) + + +def test_get_out_of_order_raises(tmp_path): + rng = np.random.default_rng(2) + with open_store("host", tmp_path) as store: + fill(store, rng) + store.prefetch(1, [(0, 3), (3, 6)]) + + with pytest.raises(RuntimeError, match="out of order"): + store.get(1, 3, 6) + + +def test_disk_tier_files_are_removed(tmp_path): + rng = np.random.default_rng(3) + with open_store("disk", tmp_path) as store: + fill(store, rng) + (scratch,) = tmp_path.iterdir() + store.prefetch(3, [(0, CHI)]) + store.get(3, 0, CHI) + store.drop(3) + assert sorted(p.name for p in scratch.iterdir()) == [ + "env-0001.bin", + "env-0002.bin", + ] + assert list(tmp_path.iterdir()) == [] + + +def fill_then_fail(tmp_path): + with open_store("disk", tmp_path) as store: + fill(store, np.random.default_rng(4)) + raise KeyError + + +def test_scratch_is_removed_after_an_error(tmp_path): + with pytest.raises(KeyError): + fill_then_fail(tmp_path) + + assert list(tmp_path.iterdir()) == [] + + +def fill_then_read(tmp_path): + with open_store("disk", tmp_path) as store: + fill(store, np.random.default_rng(5)) + store.get(1, 0, 3) + + +def test_writer_errors_are_raised(monkeypatch, tmp_path): + def full_disk(*_args): + raise OSError(errno.ENOSPC, "No space left on device") + + monkeypatch.setattr(store_module, "_write_all", full_disk) + + with pytest.raises(OSError, match=r"env-0001\.bin failed \(No space left"): + fill_then_read(tmp_path) + + assert list(tmp_path.iterdir()) == [] diff --git a/tests/test_tensor_train.py b/tests/test_tensor_train.py index f48d598..c4966b9 100644 --- a/tests/test_tensor_train.py +++ b/tests/test_tensor_train.py @@ -4,9 +4,12 @@ import pytest from src_method._tensor_train import ( + SwappedLegs, exact_stack, normalize_stack, pad, + pad_site, + padded_shape, transpose_mpo, unpad, ) @@ -174,3 +177,45 @@ def test_exact_stack_depth_three(kind, rng): want = dense_two_site(A) @ dense_two_site(B) @ dense_two_site(last) np.testing.assert_allclose(dense_two_site(out), want, atol=1e-10) + + +# ------------------------------- +# --- padded_shape / pad_site --- +# ------------------------------- + + +@pytest.mark.parametrize("kind", ["mps", "mpo"]) +@pytest.mark.parametrize("n_sites", [2, 3, 4]) +def test_padded_shape_matches_pad(kind, n_sites, rng): + train = ( + random_mps(n_sites, 3, rng) if kind == "mps" else random_mpo(n_sites, 3, rng) + ) + last = n_sites - 1 + + shapes = [padded_shape(t.shape, kind, i, last) for i, t in enumerate(train)] + + assert shapes == [t.shape for t in pad(train)] + for i, site in enumerate(train): + np.testing.assert_array_equal(pad_site(site, kind, i, last), pad(train)[i]) + + +class LazySite: + """A site with shape and dtype but no array methods, read by ``np.asarray``.""" + + def __init__(self, data): + self.data = data + self.shape, self.dtype, self.ndim = data.shape, data.dtype, data.ndim + + def __array__(self, dtype=None, copy=None): + return np.asarray(self.data, dtype=dtype) + + +def test_transpose_mpo_of_lazy_sites(rng): + train = random_mpo(3, 2, rng, up=2, down=3) + + transposed = transpose_mpo([LazySite(site) for site in train]) + + assert all(isinstance(site, SwappedLegs) for site in transposed) + assert [t.shape for t in transposed] == [(2, 3, 2), (2, 2, 3, 2), (2, 3, 2)] + np.testing.assert_array_equal(np.asarray(transposed[1]), train[1].swapaxes(-2, -1)) + assert np.asarray(transposed[0], dtype=np.complex64).dtype == np.complex64 diff --git a/uv.lock b/uv.lock index ef2c258..f548f87 100644 --- a/uv.lock +++ b/uv.lock @@ -397,6 +397,43 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/11/d0/c177e29701cf1d3008d7d2b16b5fc626592ce13bd535f8795c5f57187e0e/cuda_pathfinder-1.5.4-py3-none-any.whl", hash = "sha256:9563d3175ce1828531acf4b94e1c1c7d67208c347ca002493e2654878b26f4b7", size = 51657, upload-time = "2026-04-27T22:42:07.712Z" }, ] +[[package]] +name = "cuda-toolkit" +version = "13.4.2" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/80/02/ba4b5eaec47fb1ae01138ab0d88c735e14e594a7c1bd8146cdd4dfe15d89/cuda_toolkit-13.4.2-py2.py3-none-any.whl", hash = "sha256:2e79d99df4f3c5b3102fa4a5a69eb4c37a06ebf06b9aef92c18078c26e8db004", size = 2693, upload-time = "2026-09-16T20:57:11.438Z" }, +] + +[package.optional-dependencies] +cublas = [ + { name = "nvidia-cublas", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, + { name = "nvidia-cuda-nvrtc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] +cudart = [ + { name = "nvidia-cuda-runtime", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] +cufft = [ + { name = "nvidia-cufft", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, + { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] +curand = [ + { name = "nvidia-curand", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] +cusolver = [ + { name = "nvidia-cublas", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, + { name = "nvidia-cusolver", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, + { name = "nvidia-cusparse", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, + { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] +cusparse = [ + { name = "nvidia-cusparse", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, + { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] +nvrtc = [ + { name = "nvidia-cuda-nvrtc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] + [[package]] name = "cupy" version = "14.2.0" @@ -408,7 +445,7 @@ dependencies = [ sdist = { url = "https://files.pythonhosted.org/packages/dc/10/059a92d40b161ea24607c9e41a75098812ca7021f2b565524f87beba15c9/cupy-14.2.0.tar.gz", hash = "sha256:cfed725e612178cfbdbf62f4a3a598b60edecbbca718c0d4580c21a16ae4a62c", size = 4057114, upload-time = "2026-08-20T02:42:12.036Z" } [[package]] -name = "cupy-cuda12x" +name = "cupy-cuda13x" version = "14.2.0" source = { registry = "https://pypi.org/simple" } dependencies = [ @@ -416,21 +453,31 @@ dependencies = [ { name = "numpy" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/00/98/ac56fb7a285e264a0f29ea71d64b5c2eacd23c1f4ed9b4a8f99b16db3881/cupy_cuda12x-14.2.0-cp311-cp311-manylinux2014_aarch64.whl", hash = "sha256:1c775069f0af34662a8d4ae90848e29afcaf4ba63762d556ff22b6011683e571", size = 148189237, upload-time = "2026-08-20T02:39:47.011Z" }, - { url = "https://files.pythonhosted.org/packages/d3/49/a83b7664151a7bdfb5d7ca7f29cef4eb5574a4cb8e1f9dfbae7fea372e4f/cupy_cuda12x-14.2.0-cp311-cp311-manylinux2014_x86_64.whl", hash = "sha256:5fe2366cc5c61a7ee4a527ce1e8951cb89092d0fb0b5830623cf114d1942c585", size = 136632369, upload-time = "2026-08-20T02:39:51.562Z" }, - { url = "https://files.pythonhosted.org/packages/a0/d0/a3f4c7b7c4d642c7c8cf8ae6128ccd70cb05592f35b89d76281456e3de00/cupy_cuda12x-14.2.0-cp311-cp311-win_amd64.whl", hash = "sha256:eceffbf02a5833c8ba1c94615da07c374284db76a60f8c8b217b0d9d2667162a", size = 98876640, upload-time = "2026-08-20T02:39:55.735Z" }, - { url = "https://files.pythonhosted.org/packages/d3/8c/5fe3f6719c2d4560c79c62ef6d9b7d6c34d145879ddc0c1a41f8153ad0a6/cupy_cuda12x-14.2.0-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:b74340aa7271f0f081f77e2e5107bac75af19b86df29213db7ada90e14428efe", size = 147490427, upload-time = "2026-08-20T02:40:00.196Z" }, - { url = "https://files.pythonhosted.org/packages/7c/5b/65124de2dbaf2e85109f611a41947e39acd6dd938751c04b4c4d7bf6fc82/cupy_cuda12x-14.2.0-cp312-cp312-manylinux2014_x86_64.whl", hash = "sha256:f82141761f2c81905d49387464ae29438887956d99063381c93a1d5d1b7d32e8", size = 136111459, upload-time = "2026-08-20T02:40:04.909Z" }, - { url = "https://files.pythonhosted.org/packages/e9/18/ddea819204701024bef7fa748730702245d803847c841b737723b94fd091/cupy_cuda12x-14.2.0-cp312-cp312-win_amd64.whl", hash = "sha256:c9571d3b5f2e65758137e210f7fb3c3b34767f0af6b6ca04035a244b6141ee12", size = 98780018, upload-time = "2026-08-20T02:40:09.465Z" }, - { url = "https://files.pythonhosted.org/packages/7a/4f/dce7be227a845943d14baef3b58be49c74a465e5d9251f38840b5b1fd89a/cupy_cuda12x-14.2.0-cp313-cp313-manylinux2014_aarch64.whl", hash = "sha256:cfe673f73599ee0b9c2c9de5c0bb2395d98c9238c24deafa2ddcc69cacbd6af6", size = 147070561, upload-time = "2026-08-20T02:40:14.556Z" }, - { url = "https://files.pythonhosted.org/packages/c9/02/520f7b9f92114b4df7d88aa77c36db0d556caf76a362537687e3a2e42833/cupy_cuda12x-14.2.0-cp313-cp313-manylinux2014_x86_64.whl", hash = "sha256:efc1da23505e88d9834a3ddd3c00352c34e58e301f512d9dd593cc4bfbbdf7dc", size = 135714387, upload-time = "2026-08-20T02:40:19.077Z" }, - { url = "https://files.pythonhosted.org/packages/29/94/2dfb330afc6756ab9a8d16e955c0458e82e769930eab01e6c491e411363d/cupy_cuda12x-14.2.0-cp313-cp313-win_amd64.whl", hash = "sha256:dcea9f2b1887ac631a9275a61577e09d1eea26bf5f95491501c3b7528cebc592", size = 98760068, upload-time = "2026-08-20T02:40:23.43Z" }, - { url = "https://files.pythonhosted.org/packages/7e/d3/f6639af54f5872d1ef0c523601c7fe76d28783e71a3e8533e096c9ca1d43/cupy_cuda12x-14.2.0-cp314-cp314-manylinux2014_aarch64.whl", hash = "sha256:ed317136439af4780f217eda0b82f25180084eb16c44854e1bc9e055f96fd429", size = 146942115, upload-time = "2026-08-20T02:40:28.484Z" }, - { url = "https://files.pythonhosted.org/packages/04/5e/e6134253265fefc0a35356adcebc4e3ffa81f6c9a2a74f8f9e2de32b3018/cupy_cuda12x-14.2.0-cp314-cp314-manylinux2014_x86_64.whl", hash = "sha256:db802e4b9a85ed84fd3e84790586c06e808ee45e0214cd4e80734c09fcf93073", size = 135091580, upload-time = "2026-08-20T02:40:33.351Z" }, - { url = "https://files.pythonhosted.org/packages/0a/98/4d3215440b7a0d8661295050653760b57f32c933f1ef1c81841b7329209e/cupy_cuda12x-14.2.0-cp314-cp314-win_amd64.whl", hash = "sha256:5f08fc1d651d2446c1d18ad94f1a710224fab36d46634d4aa356423926964591", size = 99344061, upload-time = "2026-08-20T02:40:37.459Z" }, - { url = "https://files.pythonhosted.org/packages/3a/e9/8ed4adeb8c64f188b9ea6fba3be62fb7999584308bdf7ec6c5e17f77b99c/cupy_cuda12x-14.2.0-cp314-cp314t-manylinux2014_aarch64.whl", hash = "sha256:9dd33f9cfc7aefbd935879bf50e95db539721a0702bdb05be3c74bd46a85ba29", size = 147438535, upload-time = "2026-08-20T02:40:42.11Z" }, - { url = "https://files.pythonhosted.org/packages/5a/c9/73227968a5b01ac31eaf1d5c58b4318e4b83654ed6dac3c310c7b2075c36/cupy_cuda12x-14.2.0-cp314-cp314t-manylinux2014_x86_64.whl", hash = "sha256:8cbbd48c9cfd6b78d0a833ebbafda3e1b057c38d6acc3c6e54de0735a7364e27", size = 135455541, upload-time = "2026-08-20T02:40:46.804Z" }, - { url = "https://files.pythonhosted.org/packages/2e/3d/26127dd01e08ed645a70b4084ef6dde93e6a75b0a84fddc3ac6b11b05bf7/cupy_cuda12x-14.2.0-cp314-cp314t-win_amd64.whl", hash = "sha256:d14b651ed835079f8a5e273936e02eda690be7d30f2658e5f48f328322fd9d7b", size = 100242434, upload-time = "2026-08-20T02:40:51.04Z" }, + { url = "https://files.pythonhosted.org/packages/8d/3c/4716e285f8459a5707ef39d84e9a7547bf16f87c583ac4789c0763d6d986/cupy_cuda13x-14.2.0-cp311-cp311-manylinux2014_aarch64.whl", hash = "sha256:a22e2f17c97940cd566f351f186d8fe04285280d052ca1512ff223febbe9979c", size = 74074375, upload-time = "2026-08-20T02:41:04.327Z" }, + { url = "https://files.pythonhosted.org/packages/ac/cb/41b6d8f1880e6f3128f8084b7622b1a3ca2d615ab77af9ef313b4aafed3d/cupy_cuda13x-14.2.0-cp311-cp311-manylinux2014_x86_64.whl", hash = "sha256:f94843e54500a9c8c169135f79572bfaba946efe4f7178e834c227280c1b0c12", size = 69969199, upload-time = "2026-08-20T02:41:07.842Z" }, + { url = "https://files.pythonhosted.org/packages/f8/eb/25b8c1a415ade6479ed684c5cbff1e6f5002a92cbf26ef826df56a40cf69/cupy_cuda13x-14.2.0-cp311-cp311-win_amd64.whl", hash = "sha256:3893acd8d4dc78d113736045eaa306e496ee7bd8253bd959513ba28c91ed396d", size = 36082966, upload-time = "2026-08-20T02:41:10.723Z" }, + { url = "https://files.pythonhosted.org/packages/10/72/ad885cc21a21208b94044b34f50a840394b2e9dbb6111ed43959366aa6eb/cupy_cuda13x-14.2.0-cp311-cp311-win_arm64.whl", hash = "sha256:e0ad42a1ef35549282e2087a7c1e6ec8b7aa624a95e8fc0a60da777ca87e71de", size = 35136070, upload-time = "2026-09-01T06:41:38.284Z" }, + { url = "https://files.pythonhosted.org/packages/dd/c8/68e07ff959f6a9f93169e37a64a4c24223487102b8625191e4a566cab465/cupy_cuda13x-14.2.0-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:0a9b6538079c5151dcd99cf44b94cc3f1242315b7de0498b4619cef9b2a7c2c8", size = 73382974, upload-time = "2026-08-20T02:41:13.648Z" }, + { url = "https://files.pythonhosted.org/packages/7e/bf/c4c25e97552320451afdcd459d4b6e0b715459594e998b23051c06db64c3/cupy_cuda13x-14.2.0-cp312-cp312-manylinux2014_x86_64.whl", hash = "sha256:9aafadbe29db3d4285026044aafeb28f2a3de4bd0df70718ec6892b7d27a4ded", size = 69439247, upload-time = "2026-08-20T02:41:17.375Z" }, + { url = "https://files.pythonhosted.org/packages/5e/88/0a4d2d075485c4a82bfcd4725300ff87d17de36dce27b36f720d286f0f70/cupy_cuda13x-14.2.0-cp312-cp312-win_amd64.whl", hash = "sha256:9e8bc8a7693c5c212cd1450d728f47e1f05aa07ae003de30415893d0daf7d243", size = 35986921, upload-time = "2026-08-20T02:41:20.253Z" }, + { url = "https://files.pythonhosted.org/packages/a4/ab/8a756f553397b8509906d4a1d8083d92d3d21cbde115504e7a5db97de61f/cupy_cuda13x-14.2.0-cp312-cp312-win_arm64.whl", hash = "sha256:0dae6d788752c987c1cd87adb903640620820feb895339e8bc51a01fba0fb804", size = 35049803, upload-time = "2026-09-01T06:41:43.442Z" }, + { url = "https://files.pythonhosted.org/packages/3c/02/3cb39774c71a16275c50b5d1e92ddcefa948ff1e69f464249b70a79b0abe/cupy_cuda13x-14.2.0-cp313-cp313-manylinux2014_aarch64.whl", hash = "sha256:449af7f608dc70ce64e8c45cf5a6ea0bbef51623e7c806b1a43445ae3c76e1e0", size = 72969322, upload-time = "2026-08-20T02:41:23.311Z" }, + { url = "https://files.pythonhosted.org/packages/4e/27/216370ae8c39fa9835645be3c27513ab72f52cf4f51faf58d799fba981b6/cupy_cuda13x-14.2.0-cp313-cp313-manylinux2014_x86_64.whl", hash = "sha256:5b79a606f6639d9bed74ef7072f36756a95a804c86d7adb367b39307ac771600", size = 69048238, upload-time = "2026-08-20T02:41:26.831Z" }, + { url = "https://files.pythonhosted.org/packages/5d/75/9df3c545baad4ff2e384db1ff6c8005fa4040492c4ee5672ca8b8711d29d/cupy_cuda13x-14.2.0-cp313-cp313-win_amd64.whl", hash = "sha256:f9a2a55abf889f6df68b5f62cd5f3ec4a2f305ac57ed09becd52e16201edcd86", size = 35965490, upload-time = "2026-08-20T02:41:29.943Z" }, + { url = "https://files.pythonhosted.org/packages/5c/3d/ee85ce39c2f32e8561c81c507df458a914b45d9a0eb1a54e5396c2cfd632/cupy_cuda13x-14.2.0-cp313-cp313-win_arm64.whl", hash = "sha256:5ddaf38291f60ec20e8c7889b28e339bf17ad1b4df768b8bb5e6c9f9e445efbb", size = 35029087, upload-time = "2026-09-01T06:41:48.593Z" }, + { url = "https://files.pythonhosted.org/packages/e5/70/8c80c9e7010193afed8ccb7110d6173aa23bf3a6ce2fc2450244b0fa50a1/cupy_cuda13x-14.2.0-cp314-cp314-manylinux2014_aarch64.whl", hash = "sha256:7ac57c4e16c62f265b5ecc2b45d6bbafe9dc70a99b849ed6cfb8d16f32204589", size = 72840193, upload-time = "2026-08-20T02:41:33.146Z" }, + { url = "https://files.pythonhosted.org/packages/c2/b8/4f4c4f34fc31ab8d136ed965a919505297974629d9852c50e15bfd616281/cupy_cuda13x-14.2.0-cp314-cp314-manylinux2014_x86_64.whl", hash = "sha256:ff0bdebd1b43c0c6db53095784c787c4e4eae671356cf521c4b7482ed78a1a7e", size = 68422828, upload-time = "2026-08-20T02:41:36.527Z" }, + { url = "https://files.pythonhosted.org/packages/7d/5c/58941ece5ed0b2cc8304f8ba81b5e8334759d409306475d53c52ac99809b/cupy_cuda13x-14.2.0-cp314-cp314-win_amd64.whl", hash = "sha256:b99a8bf9d5391954c0c08675a8cb39219923fe2ee249faa3546ebbf75783a360", size = 36110882, upload-time = "2026-08-20T02:41:39.837Z" }, + { url = "https://files.pythonhosted.org/packages/e2/c5/cc7d1a9b764b3f828e09b48a73d0fdeaad15ecd2c7e50507f6d5fcf5fd1a/cupy_cuda13x-14.2.0-cp314-cp314-win_arm64.whl", hash = "sha256:c61c59ba665948c0a855926997ee96026420b628816b2e406998ec797caddb22", size = 35184806, upload-time = "2026-09-01T06:41:58.317Z" }, + { url = "https://files.pythonhosted.org/packages/6f/54/a582ac70e9e75d3fad7ae073051bcd50dd128a1f66dd275291dc5f48d997/cupy_cuda13x-14.2.0-cp314-cp314t-manylinux2014_aarch64.whl", hash = "sha256:e8bb453551deced08ec5c4576da61ebf1e2b4a45aeff48e09f47ec332f7cad3d", size = 73336756, upload-time = "2026-08-20T02:41:42.718Z" }, + { url = "https://files.pythonhosted.org/packages/f9/aa/222151ba5061d7a01efe4009efad9f9a7164e15bb2a494d9aa86efc97d4f/cupy_cuda13x-14.2.0-cp314-cp314t-manylinux2014_x86_64.whl", hash = "sha256:313500c5d415f65d388e84a9d2eb7f07f97c23c85dc5785f986e75d46d9c9693", size = 68792663, upload-time = "2026-08-20T02:41:46.079Z" }, + { url = "https://files.pythonhosted.org/packages/d4/20/13dd37dbeaec879c8348fde2bfca67d96daa56a6e63e4a5083ff61660e53/cupy_cuda13x-14.2.0-cp314-cp314t-win_amd64.whl", hash = "sha256:c517fc55502a7e3d1a8b9310abad95cec61dded5a38d328fe0220bcf7b036268", size = 37004845, upload-time = "2026-08-20T02:41:48.963Z" }, + { url = "https://files.pythonhosted.org/packages/d1/5b/b1367f79abd90a259dbb7f0603dd162c03c3a76ec1d02c2527b0c79d46ba/cupy_cuda13x-14.2.0-cp314-cp314t-win_arm64.whl", hash = "sha256:83885db963eb0b95a0757b4907a37115a1d2aebfa9cf43b99ebac8db35713d37", size = 35485280, upload-time = "2026-09-01T06:41:53.618Z" }, +] + +[package.optional-dependencies] +ctk = [ + { name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "curand", "cusolver", "cusparse", "nvrtc"] }, ] [[package]] @@ -1404,74 +1451,94 @@ wheels = [ ] [[package]] -name = "nvidia-cublas-cu12" -version = "12.9.2.10" +name = "nvidia-cublas" +version = "13.8.0.4" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cuda-nvrtc-cu12" }, + { name = "nvidia-cuda-nvrtc" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/c9/9e/0c73b14c14cd32e8b3a028133371f4bf1d9d14b811c515c5fd1201a3c38d/nvidia_cublas-13.8.0.4-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:e22b25be18f8b7d267dbbe2a8ec78e2c69c05762536d299dd0afa2c847115740", size = 552729449, upload-time = "2026-09-16T20:40:37.357Z" }, + { url = "https://files.pythonhosted.org/packages/7a/38/bdd540bf511d2c9b6f9efc71a81c60cb88e295be0b9312b61d19bbed2212/nvidia_cublas-13.8.0.4-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:9f17797dfcc048694461f4e47de17d2e3c25adf172ef723d2db0a07cd8744b89", size = 439317144, upload-time = "2026-09-16T20:41:11.569Z" }, + { url = "https://files.pythonhosted.org/packages/a3/df/f1246959833e2c437db8be3e5b477f66b87f8817821ed40de6c7561c9a36/nvidia_cublas-13.8.0.4-py3-none-win_amd64.whl", hash = "sha256:8c5494423bb8a46822cb6b0cb95d7fa4be2d7b96a31155dff083839ec8297910", size = 423266897, upload-time = "2026-09-16T20:48:16.789Z" }, + { url = "https://files.pythonhosted.org/packages/96/38/723c97681824250b261c91761de3f1a1899111a9b1dda65eec5b2daa1d58/nvidia_cublas-13.8.0.4-py3-none-win_arm64.whl", hash = "sha256:a2ffda7a27d8315e6b75c1ced526d5a68c7a2a50adfe866235401c9592d3c9e4", size = 153856980, upload-time = "2026-09-16T20:54:58.249Z" }, ] + +[[package]] +name = "nvidia-cuda-nvrtc" +version = "13.4.92" +source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f7/a2/c96163a0fff1839c0c9548bbdeae7b853b867009e33b9b9264adc238b1cf/nvidia_cublas_cu12-12.9.2.10-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:5572131a59c3eebeeb1c4c8144f772d49372c20124916e072a0e3fc30df421d5", size = 575012079, upload-time = "2026-04-08T18:51:47.303Z" }, - { url = "https://files.pythonhosted.org/packages/cb/c0/0a517bfe63ccd3b92eb254d264e28fca3c7cab75d07daea315250fb1bf73/nvidia_cublas_cu12-12.9.2.10-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:e4f53a8ca8c5d6e8c492d0d0a3d565ecb59a751b19cfdaa4f6da0ab2104c1702", size = 581240110, upload-time = "2026-04-08T18:52:31.532Z" }, - { url = "https://files.pythonhosted.org/packages/20/e2/fc9a0e985249d873150276d5afb02e39a66817fedbf1a385724393e505ed/nvidia_cublas_cu12-12.9.2.10-py3-none-win_amd64.whl", hash = "sha256:623f43027d40d44ceadf0043f002bd25cf353e8f13ce90b9a87057019f560661", size = 553162896, upload-time = "2026-04-08T18:53:10.035Z" }, + { url = "https://files.pythonhosted.org/packages/56/9c/1342ebb460ce2afd014ec5a002adc99108e3c53a6b70f399cf64dd1f267d/nvidia_cuda_nvrtc-13.4.92-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:5ce8c97b00b232c4f50c8c4b5a3b68cafee08bdb82ea86f2052ff01d03194f4a", size = 53301733, upload-time = "2026-09-16T20:39:06.686Z" }, + { url = "https://files.pythonhosted.org/packages/4f/73/cf76dc0083d68bd1e29309bc51fc95ada28d4b64d341e945eb3f8ec8edec/nvidia_cuda_nvrtc-13.4.92-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:24b9f5eccc6a5a19779038cf468aecb7cecfa8716269beaef5640bd989d21c28", size = 50927348, upload-time = "2026-09-16T20:38:55.548Z" }, + { url = "https://files.pythonhosted.org/packages/64/5b/aa91896f64444eff1dd7f189f4db05046c7f0ba08c0f531283c1c1aac581/nvidia_cuda_nvrtc-13.4.92-py3-none-win_amd64.whl", hash = "sha256:6af7ac5372920f6a7a560d0699348560fe44cb52c1d722d6afd3119f8af948c4", size = 46972283, upload-time = "2026-09-16T20:47:13.295Z" }, + { url = "https://files.pythonhosted.org/packages/25/22/af15f5ef51ba0f90d5439487b8efe482d7b2cdafa8ff6a01e13791622b98/nvidia_cuda_nvrtc-13.4.92-py3-none-win_arm64.whl", hash = "sha256:1620066e967e93119d67338628935cbb1196b53a1474b07f14b7b77aba477284", size = 42709926, upload-time = "2026-09-16T20:54:16.396Z" }, ] [[package]] -name = "nvidia-cuda-nvrtc-cu12" -version = "12.9.86" +name = "nvidia-cuda-runtime" +version = "13.4.92" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b8/85/e4af82cc9202023862090bfca4ea827d533329e925c758f0cde964cb54b7/nvidia_cuda_nvrtc_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:210cf05005a447e29214e9ce50851e83fc5f4358df8b453155d5e1918094dcb4", size = 89568129, upload-time = "2025-06-05T20:02:41.973Z" }, - { url = "https://files.pythonhosted.org/packages/64/eb/c2295044b8f3b3b08860e2f6a912b702fc92568a167259df5dddb78f325e/nvidia_cuda_nvrtc_cu12-12.9.86-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:096d4de6bda726415dfaf3198d4f5c522b8e70139c97feef5cd2ca6d4cd9cead", size = 44528905, upload-time = "2025-06-05T20:02:29.754Z" }, - { url = "https://files.pythonhosted.org/packages/52/de/823919be3b9d0ccbf1f784035423c5f18f4267fb0123558d58b813c6ec86/nvidia_cuda_nvrtc_cu12-12.9.86-py3-none-win_amd64.whl", hash = "sha256:72972ebdcf504d69462d3bcd67e7b81edd25d0fb85a2c46d3ea3517666636349", size = 76408187, upload-time = "2025-06-05T20:12:27.819Z" }, + { url = "https://files.pythonhosted.org/packages/51/3b/2c6e9d88efa7b3585572760bca5e7da83fa9e06098a9386c7b29c81c46b2/nvidia_cuda_runtime-13.4.92-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:bef071788589550ab02846fcbc732a2b88c2b16f2cd8a0f45d98e68e8c5fb0c8", size = 2509438, upload-time = "2026-09-16T20:37:16.36Z" }, + { url = "https://files.pythonhosted.org/packages/98/8a/3431271f6344874b8f1ac03f16b3d679c91493f8da63f716160403e6d0a0/nvidia_cuda_runtime-13.4.92-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:9641f797da20ce1dd8e779b6e96d08cf9ba564cec8e8225458811ee26423f3a5", size = 2494438, upload-time = "2026-09-16T20:37:21.89Z" }, + { url = "https://files.pythonhosted.org/packages/86/00/d5436004268f049214193659ebc36550b5ef3925c3d13b4cc980e13be6f5/nvidia_cuda_runtime-13.4.92-py3-none-win_amd64.whl", hash = "sha256:08dca5e4aba480c2fd5b55075c0fa71b84ef9dcf0521f2d58baa14a803a7311c", size = 2778543, upload-time = "2026-09-16T20:46:42.459Z" }, + { url = "https://files.pythonhosted.org/packages/b5/bc/9a141b73d65b9f09e548fcb9b7f4415da22f35e0fe78c60e3c0a9b458efd/nvidia_cuda_runtime-13.4.92-py3-none-win_arm64.whl", hash = "sha256:43972819798ca06ad6354f6cbdb82a0a113f88b94e9ea90256d660027b0b3b60", size = 2770284, upload-time = "2026-09-16T20:53:45.718Z" }, ] [[package]] -name = "nvidia-cuda-runtime-cu12" -version = "12.9.79" +name = "nvidia-cufft" +version = "12.4.0.43" source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink" }, +] wheels = [ - { url = "https://files.pythonhosted.org/packages/bc/e0/0279bd94539fda525e0c8538db29b72a5a8495b0c12173113471d28bce78/nvidia_cuda_runtime_cu12-12.9.79-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:83469a846206f2a733db0c42e223589ab62fd2fabac4432d2f8802de4bded0a4", size = 3515012, upload-time = "2025-06-05T20:00:35.519Z" }, - { url = "https://files.pythonhosted.org/packages/bc/46/a92db19b8309581092a3add7e6fceb4c301a3fd233969856a8cbf042cd3c/nvidia_cuda_runtime_cu12-12.9.79-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:25bba2dfb01d48a9b59ca474a1ac43c6ebf7011f1b0b8cc44f54eb6ac48a96c3", size = 3493179, upload-time = "2025-06-05T20:00:53.735Z" }, - { url = "https://files.pythonhosted.org/packages/59/df/e7c3a360be4f7b93cee39271b792669baeb3846c58a4df6dfcf187a7ffab/nvidia_cuda_runtime_cu12-12.9.79-py3-none-win_amd64.whl", hash = "sha256:8e018af8fa02363876860388bd10ccb89eb9ab8fb0aa749aaf58430a9f7c4891", size = 3591604, upload-time = "2025-06-05T20:11:17.036Z" }, + { url = "https://files.pythonhosted.org/packages/41/f1/bb8cc4e1ffee345a51fe4cf656352bc705274431dcbd6c23ed9f00585ec7/nvidia_cufft-12.4.0.43-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:3938644a5b594e06d396e02d6e52bdd84589b3c1f3a8a127a2704211c2abc441", size = 161750525, upload-time = "2026-09-16T20:41:56.892Z" }, + { url = "https://files.pythonhosted.org/packages/76/bf/3fea3d1c6262bded26ae00e3106432d63235954965d7785f5051ac146651/nvidia_cufft-12.4.0.43-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:0e8385013596b112d29c9ce8c63dc575b308d77636c7169104e18714f03961a8", size = 161750089, upload-time = "2026-09-16T20:42:15.469Z" }, + { url = "https://files.pythonhosted.org/packages/7b/cc/be7fe31058127336a66c88414a2ecc6beecf680baf85994a43eaf292b202/nvidia_cufft-12.4.0.43-py3-none-win_amd64.whl", hash = "sha256:4ff7075f2d0b5f69291f70938d37a86ec632cbe5747184c74ba1f50f17accacc", size = 160953139, upload-time = "2026-09-16T20:48:43.428Z" }, + { url = "https://files.pythonhosted.org/packages/0b/ee/9217a02eae5b0b183b62b39b65f1e64ee8ffb6d6da67a79fd7a2f295ed3a/nvidia_cufft-12.4.0.43-py3-none-win_arm64.whl", hash = "sha256:4e8d551542bd661aef431422acddc5f06f793ac9b1dc7af53de63cbfde5781f9", size = 161850334, upload-time = "2026-09-16T20:55:16.698Z" }, ] [[package]] -name = "nvidia-curand-cu12" -version = "10.3.10.19" +name = "nvidia-curand" +version = "10.4.4.72" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/14/1c/2a45afc614d99558d4a773fa740d8bb5471c8398eeed925fc0fcba020173/nvidia_curand_cu12-10.3.10.19-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:de663377feb1697e1d30ed587b07d5721fdd6d2015c738d7528a6002a6134d37", size = 68292066, upload-time = "2025-05-01T19:39:13.595Z" }, - { url = "https://files.pythonhosted.org/packages/31/44/193a0e171750ca9f8320626e8a1f2381e4077a65e69e2fb9708bd479e34a/nvidia_curand_cu12-10.3.10.19-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:49b274db4780d421bd2ccd362e1415c13887c53c214f0d4b761752b8f9f6aa1e", size = 68295626, upload-time = "2025-05-01T19:39:38.885Z" }, - { url = "https://files.pythonhosted.org/packages/e5/98/1bd66fd09cbe1a5920cb36ba87029d511db7cca93979e635fd431ad3b6c0/nvidia_curand_cu12-10.3.10.19-py3-none-win_amd64.whl", hash = "sha256:e8129e6ac40dc123bd948e33d3e11b4aa617d87a583fa2f21b3210e90c743cde", size = 68774847, upload-time = "2025-05-01T19:48:52.93Z" }, + { url = "https://files.pythonhosted.org/packages/e8/9d/42faff77e90e5498eecb9b099894d5deea744b3d6b35e700dafd0381e07a/nvidia_curand-10.4.4.72-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:53bef256d4362eb3d70c9c048a8d31bcf003f9dcade17720c73d3a494a7dcaa3", size = 63771251, upload-time = "2026-09-16T20:42:40.644Z" }, + { url = "https://files.pythonhosted.org/packages/07/73/3ee8e5b4cb891401e603ffd3a59b35c6afe785fd2de123afe7c7029603dc/nvidia_curand-10.4.4.72-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:25c3457ae7a224fdd484dab90b0fc5dc0e842fab5db3012afa4a5bd2af4eb7e5", size = 61498332, upload-time = "2026-09-16T20:42:51.685Z" }, + { url = "https://files.pythonhosted.org/packages/e2/d4/f59ab342b82ff05ff27d674acd470ba625d788bf71338f7da6c89669d6b9/nvidia_curand-10.4.4.72-py3-none-win_amd64.whl", hash = "sha256:e0bce83e083ef25976ee74f59e8f067c15149a74538f1be6c2462e286e7c9c68", size = 56731059, upload-time = "2026-09-16T20:48:53.631Z" }, + { url = "https://files.pythonhosted.org/packages/8f/60/6734d25ee55e688a88d4ea5461d71f2cd9d8f0e8da50d5d10f3e6408f855/nvidia_curand-10.4.4.72-py3-none-win_arm64.whl", hash = "sha256:4635b2c8a727f51b585614f590509cb5c50fb2fc55fa9350f01c22c7f8c0fcdf", size = 65140169, upload-time = "2026-09-16T20:55:30.897Z" }, ] [[package]] -name = "nvidia-cusolver-cu12" -version = "11.7.5.82" +name = "nvidia-cusolver" +version = "12.3.4.7" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas-cu12" }, - { name = "nvidia-cusparse-cu12" }, - { name = "nvidia-nvjitlink-cu12" }, + { name = "nvidia-cublas" }, + { name = "nvidia-cusparse" }, + { name = "nvidia-nvjitlink" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/03/99/686ff9bf3a82a531c62b1a5c614476e8dfa24a9d89067aeedf3592ee4538/nvidia_cusolver_cu12-11.7.5.82-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:62efa83e4ace59a4c734d052bb72158e888aa7b770e1a5f601682f16fe5b4fd2", size = 337869834, upload-time = "2025-06-05T20:06:53.125Z" }, - { url = "https://files.pythonhosted.org/packages/33/40/79b0c64d44d6c166c0964ec1d803d067f4a145cca23e23925fd351d0e642/nvidia_cusolver_cu12-11.7.5.82-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:15da72d1340d29b5b3cf3fd100e3cd53421dde36002eda6ed93811af63c40d88", size = 338117415, upload-time = "2025-06-05T20:07:16.809Z" }, - { url = "https://files.pythonhosted.org/packages/32/5d/feb7f86b809f89b14193beffebe24cf2e4bf7af08372ab8cdd34d19a65a0/nvidia_cusolver_cu12-11.7.5.82-py3-none-win_amd64.whl", hash = "sha256:77666337237716783c6269a658dea310195cddbd80a5b2919b1ba8735cec8efd", size = 326215953, upload-time = "2025-06-05T20:14:41.76Z" }, + { url = "https://files.pythonhosted.org/packages/77/f1/8c8a667c59cb4ea0dd1543ed99c6552d818e2a4bb922d0ee98f0a49de207/nvidia_cusolver-12.3.4.7-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:4a38d88a1ea3f7b656e52001caf943b625822a89df3a8b63d895f56ef777b1a4", size = 291070656, upload-time = "2026-09-16T20:43:15.438Z" }, + { url = "https://files.pythonhosted.org/packages/3f/43/e1a29568cc9989c95bcf31519f69baccfaebdf5ef5e624773726cc9740e2/nvidia_cusolver-12.3.4.7-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:225dd543c7b93ca22e62a35b5f4d8b4f9caa515733793925c05cdcf58da5cf02", size = 246439468, upload-time = "2026-09-16T20:43:34.367Z" }, + { url = "https://files.pythonhosted.org/packages/65/0b/1d751c92e1429ef3261021eb4b6cf27076b547d7512c6dc4b9aebeb9c49a/nvidia_cusolver-12.3.4.7-py3-none-win_amd64.whl", hash = "sha256:7ed56898cd98abe36d8727eaf1408622c67557f2acb3baca5a27e24e595a1ccb", size = 261421900, upload-time = "2026-09-16T20:49:15.844Z" }, + { url = "https://files.pythonhosted.org/packages/51/e7/e001caaf3101d01d80583753253a041b47db91e67fe71fb6320ff7afd7c7/nvidia_cusolver-12.3.4.7-py3-none-win_arm64.whl", hash = "sha256:87891ab21de591da154070bd58d5fd3ab52ccaeb900f741a7ead0c8a859e44ac", size = 71318077, upload-time = "2026-09-16T20:55:41.128Z" }, ] [[package]] -name = "nvidia-cusparse-cu12" -version = "12.5.10.65" +name = "nvidia-cusparse" +version = "12.8.6.72" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink-cu12" }, + { name = "nvidia-nvjitlink" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/5e/6f/8710fbd17cdd1d0fc3fea7d36d5b65ce1933611c31e1861da330206b253a/nvidia_cusparse_cu12-12.5.10.65-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:221c73e7482dd93eda44e65ce567c031c07e2f93f6fa0ecd3ba876a195023e83", size = 366359408, upload-time = "2025-06-05T20:07:42.501Z" }, - { url = "https://files.pythonhosted.org/packages/12/46/b0fd4b04f86577921feb97d8e2cf028afe04f614d17fb5013de9282c9216/nvidia_cusparse_cu12-12.5.10.65-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:73060ce019ac064a057267c585bf1fd5a353734151f87472ff02b2c5c9984e78", size = 366465088, upload-time = "2025-06-05T20:08:20.413Z" }, - { url = "https://files.pythonhosted.org/packages/73/ef/063500c25670fbd1cbb0cd3eb7c8a061585b53adb4dd8bf3492bb49b0df3/nvidia_cusparse_cu12-12.5.10.65-py3-none-win_amd64.whl", hash = "sha256:9e487468a22a1eaf1fbd1d2035936a905feb79c4ce5c2f67626764ee4f90227c", size = 362504719, upload-time = "2025-06-05T20:15:17.947Z" }, + { url = "https://files.pythonhosted.org/packages/d1/2e/bf04fbd787d6b227da62abbe3674de5d89460dacebbf5cc97b5717b89a49/nvidia_cusparse-12.8.6.72-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:c3917c86fd419cbd42229b5cccb8bbe5200ed5b7dec3aa86d28e25c90a256e1b", size = 188672939, upload-time = "2026-09-16T20:43:52.751Z" }, + { url = "https://files.pythonhosted.org/packages/80/cb/409df72613bf1c6cd9b545a257ab5489ecfd98c88cdf7ee9d7fc434e94a1/nvidia_cusparse-12.8.6.72-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a739f6ff51ea2a8a2990b267b9b4d3063a5d94890efb8ccf393bfcab9dd487aa", size = 170468146, upload-time = "2026-09-16T20:44:11.113Z" }, + { url = "https://files.pythonhosted.org/packages/85/f7/e0d0edefab227c9967df12484fd50e5586b9af86439f38354779c7086e8b/nvidia_cusparse-12.8.6.72-py3-none-win_amd64.whl", hash = "sha256:013d3f83316431dd3338cf15141dd5420b24dffc6663fc083c053c8409fa3cb7", size = 168202645, upload-time = "2026-09-16T20:49:32.639Z" }, + { url = "https://files.pythonhosted.org/packages/23/ff/bedb3859cedf60ba14c761181f91bcfe926fb072bf686deb5f93c6cb7b10/nvidia_cusparse-12.8.6.72-py3-none-win_arm64.whl", hash = "sha256:7efe07b54c505f3eeec0469078c992a31b15ab716d40f07a8109f7cbcc4fd7fe", size = 184995442, upload-time = "2026-09-16T20:55:58.013Z" }, ] [[package]] @@ -1484,13 +1551,14 @@ wheels = [ ] [[package]] -name = "nvidia-nvjitlink-cu12" -version = "12.9.86" +name = "nvidia-nvjitlink" +version = "13.4.92" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/46/0c/c75bbfb967457a0b7670b8ad267bfc4fffdf341c074e0a80db06c24ccfd4/nvidia_nvjitlink_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:e3f1171dbdc83c5932a45f0f4c99180a70de9bd2718c1ab77d14104f6d7147f9", size = 39748338, upload-time = "2025-06-05T20:10:25.613Z" }, - { url = "https://files.pythonhosted.org/packages/97/bc/2dcba8e70cf3115b400fef54f213bcd6715a3195eba000f8330f11e40c45/nvidia_nvjitlink_cu12-12.9.86-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:994a05ef08ef4b0b299829cde613a424382aff7efb08a7172c1fa616cc3af2ca", size = 39514880, upload-time = "2025-06-05T20:10:04.89Z" }, - { url = "https://files.pythonhosted.org/packages/dd/7e/2eecb277d8a98184d881fb98a738363fd4f14577a4d2d7f8264266e82623/nvidia_nvjitlink_cu12-12.9.86-py3-none-win_amd64.whl", hash = "sha256:cc6fcec260ca843c10e34c936921a1c426b351753587fdd638e8cff7b16bb9db", size = 35584936, upload-time = "2025-06-05T20:16:08.525Z" }, + { url = "https://files.pythonhosted.org/packages/1d/6b/eef7a9e32872b8f41e145bf10cddc9af26e153c338852811fe9a9baddf9e/nvidia_nvjitlink-13.4.92-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:e0391f24ed94ec879b84e3da4d4ec320c879aff681f2c7a638462f7199284323", size = 42452378, upload-time = "2026-09-16T20:45:29.042Z" }, + { url = "https://files.pythonhosted.org/packages/1f/a8/1cbd4014898af8b419e69b0d7dbc63da2121ee92d92b47d59f4fe9075349/nvidia_nvjitlink-13.4.92-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:25f74fad0d654271c921ac4dca614bd6258bc21791242fc7b2289dad7ae9c099", size = 40420120, upload-time = "2026-09-16T20:45:19.163Z" }, + { url = "https://files.pythonhosted.org/packages/72/7c/c44bd7277afafed9de262a23baec4dc8dcf7ecbb1b9ecb228821442d1e9b/nvidia_nvjitlink-13.4.92-py3-none-win_amd64.whl", hash = "sha256:b286f3a4f227a9363efdec263c7b91788cef1478d2b8a5fa8bab7f3e82ff82fd", size = 39053950, upload-time = "2026-09-16T20:52:54.244Z" }, + { url = "https://files.pythonhosted.org/packages/3a/8e/afaa7687fac05a507d10765f6097937e99735f835c5b262276d1e2d2cad5/nvidia_nvjitlink-13.4.92-py3-none-win_arm64.whl", hash = "sha256:9e4a7ff4f0cafa8c624917055b863dc11f5c2912ead23c462889e166f3b0e57d", size = 35981621, upload-time = "2026-09-16T20:56:30.819Z" }, ] [[package]] @@ -2372,14 +2440,7 @@ dependencies = [ [package.optional-dependencies] gpu-nvidia = [ - { name = "cupy-cuda12x" }, - { name = "nvidia-cublas-cu12" }, - { name = "nvidia-cuda-nvrtc-cu12" }, - { name = "nvidia-cuda-runtime-cu12" }, - { name = "nvidia-curand-cu12" }, - { name = "nvidia-cusolver-cu12" }, - { name = "nvidia-cusparse-cu12" }, - { name = "nvidia-nvjitlink-cu12" }, + { name = "cupy-cuda13x", extra = ["ctk"] }, ] gpu-rocm = [ { name = "cupy" }, @@ -2438,15 +2499,8 @@ test = [ [package.metadata] requires-dist = [ { name = "cupy", marker = "extra == 'gpu-rocm'", specifier = ">=14.1.1" }, - { name = "cupy-cuda12x", marker = "extra == 'gpu-nvidia'", specifier = ">=13" }, + { name = "cupy-cuda13x", extras = ["ctk"], marker = "extra == 'gpu-nvidia'", specifier = ">=14" }, { name = "numpy", specifier = ">=2.0,<3" }, - { name = "nvidia-cublas-cu12", marker = "extra == 'gpu-nvidia'" }, - { name = "nvidia-cuda-nvrtc-cu12", marker = "extra == 'gpu-nvidia'" }, - { name = "nvidia-cuda-runtime-cu12", marker = "extra == 'gpu-nvidia'" }, - { name = "nvidia-curand-cu12", marker = "extra == 'gpu-nvidia'" }, - { name = "nvidia-cusolver-cu12", marker = "extra == 'gpu-nvidia'" }, - { name = "nvidia-cusparse-cu12", marker = "extra == 'gpu-nvidia'" }, - { name = "nvidia-nvjitlink-cu12", marker = "extra == 'gpu-nvidia'" }, { name = "opt-einsum", specifier = ">=3.4.0" }, ] provides-extras = ["gpu-nvidia", "gpu-rocm"]