diff --git a/README.md b/README.md index 6c0d046..d27a385 100644 --- a/README.md +++ b/README.md @@ -125,7 +125,9 @@ out = src( On a node with several GPUs, `Resources(devices=4)` splits the sweep among four of them, in the same process. -See [Large problems](https://algorithmiq.github.io/src-method/features/large-problems) for the budgets, the scratch directory, several GPUs and how to read the logged plan. +The contractions follow opt_einsum's cheapest (`dp`) paths. `optimize` selects any other opt_einsum path optimizer, by name or as a `PathOptimizer` instance. + +See [Large problems](https://algorithmiq.github.io/src-method/features/large-problems) for the budgets, the scratch directory, several GPUs, contraction paths and how to read the logged plan. ## Installation diff --git a/benches/large/README.md b/benches/large/README.md index f2a4802..b2e8198 100644 --- a/benches/large/README.md +++ b/benches/large/README.md @@ -116,3 +116,52 @@ Speed-up and efficiency are against the single-GPU run: memory unused was not investigated. - The stall times the sweep logs include waiting for the GPU's own queue, so they cannot tell slow transfers from a GPU-bound pass. + +### Einsum paths: `auto` against `dp` + +The `optimize` argument of `src`, `apply` and `compress` selects the opt_einsum +path optimizer of the kernels, and the planner sizes batches for the same path. +It defaults to `dp` since these measurements; the results above ran with `auto`. +`bench_einsum_dp.py` compares the two: `kernels` times `SiteKernels.env`, +`.sketch` and `.project` on random arrays of the reference shapes, and `sweep` +with `distance` runs the `D_M = 1000` stack, `chi_out = 500`, under the 60 GB +budget, three times per optimizer and compares the outputs. Run on A100-SXM-64GB +(CuPy 14.2.0, NumPy 2.4.4, opt_einsum 3.4.0), complex128. + +Full sweep, median of 3 (range), every environment on the device: + +| GPUs | `auto` (s) | `dp` (s) | Speed-up | Right-to-left `auto` / `dp` (s) | Planned peak `auto` / `dp` (GB) | QR distance | +|---|---|---|---|---|---|---| +| 1 | 109.7 (109.6-130.6) | 56.2 (56.2-74.1) | 1.95x | 95.4 / 42.3 | 58.8 / 50.3 | 2.7e-13 | +| 4 | 32.4 (32.3-51.7) | 17.4 (17.4-40.1) | 1.86x | 28.6 / 13.6 | 32.3 / 12.2 per GPU | 4.0e-13 | + +- The gain is in the right-to-left pass; the left-to-right pass takes the same + time (14.3 against 13.9 s on one GPU). At these shapes `auto`'s sketch and + projection paths cost 3.9x and 2.0x `dp`'s flops. Both optimizers give the + same planned batches. +- The first sweep of each job is 18-23 s slower with either optimizer (the + first-run effect noted above), and is left out of the medians. `dp`'s path + search takes about 1 ms per shape: 0.03-0.10 s per sweep. + +Kernels at the reference shapes, `D = (4, 4, 4000, 4)`, physical legs 4, +`eta = 2000`, median of 5 after 2 warm-ups (ms per row or column): + +| Kernel | Batch | `auto` | `dp` | Flops `auto` / `dp` | +|---|---|---|---|---| +| project | 8 | 25.2 | 11.8 | 7.9e11 / 4.0e11 | +| project | 16 | 22.1 | 11.0 | 1.6e12 / 7.9e11 | +| project | 32, 64, 128 | over 50 GB | 10.6, 10.5, 10.5 | 9.7e12 / 1.6e12 at 32 | +| sketch | 96 | 10.27 | 10.65 | 4.78e12 / 4.76e12 | +| sketch | 192 | 10.14 | 10.53 | 9.56e12 / 9.51e12 | +| env | 224 | 7.17 | 7.18 | 7.4e12 / 7.4e12 | + +- The projection runs at 16-19 TFLOP/s with either optimizer, so the time follows + the flops. From 32 rows `auto` plans an 8.19e9-element intermediate (265-283 + GB), which is why it cannot leave 16 rows; `dp` needs 10-25 GB up to 128 rows. +- 384 sketch columns (75 GB) and 256 projection rows with `dp` (50.3 GB) exceed + the 50 GB limit and were not run. +- `dp` is 3.7-3.8% slower on the sketch, where its path has 0.5% fewer flops. +- The outputs of the two optimizers agree to 1.6e-14 or better. +- The reference-size sweep (over an hour) was not run. Taking these per-row and + per-column times for each site predicts its right-to-left pass at about 2060 s + with `dp`, 1.5x faster than the 3112 s measured with `auto`. diff --git a/benches/large/bench_einsum_dp.py b/benches/large/bench_einsum_dp.py new file mode 100644 index 0000000..c5c521c --- /dev/null +++ b/benches/large/bench_einsum_dp.py @@ -0,0 +1,445 @@ +"""opt_einsum ``auto`` against ``dp`` paths for the SRC kernels, on the GPU. + +Three commands: + +- ``kernels``: time `SiteKernels.env`, ``.sketch`` and ``.project`` on random arrays + of the reference shapes, with both optimizers in one process, and check that + their outputs agree; +- ``sweep``: run the SRC of a ``generate``-d stack of `bench_large` a few times with + one optimizer, logging the plan, the pass times and the device peak, and save the + last output; +- ``distance``: the QR distance between the saved ``auto`` and ``dp`` outputs. + +Each command logs its rows and appends them as JSON lines under ``OUT``. +""" + +from __future__ import annotations + +import json +import logging +import os +import re +import socket +import sys +from importlib.metadata import version +from math import prod +from pathlib import Path +from statistics import median +from time import perf_counter +from typing import TYPE_CHECKING, Any + +import cyclopts +import numpy as np +import opt_einsum as oe + +import src_method +import src_method._kernels as kernels_module +import src_method._sweep as sweep_module +from src_method import Resources, src +from src_method._plan import make_plan + +if TYPE_CHECKING: + from collections.abc import Callable + +logger = logging.getLogger(__name__) +app = cyclopts.App(help="Measure opt_einsum auto against dp for the SRC kernels.") + +HERE = Path(__file__).parent +OUT = Path("einsum_dp") +OPTIMIZERS = ("auto", "dp") +GB = 1e9 +# opt_einsum counts a real multiply-add as 2 operations; a complex one is 8 real +# flops, so complex128 rates are 4x the opt_einsum count over the time. +COMPLEX_FLOPS = 4 + + +def provenance(optimize: str | None = None) -> dict[str, Any]: + """Versions (``src_method`` carries the commit), job and GPU of this process.""" + import cupy # noqa: PLC0415 (GPU only) + + return { + "optimize": optimize, + "src_method": src_method.__version__, + "numpy": np.__version__, + "cupy": cupy.__version__, + "opt_einsum": version("opt_einsum"), + "python": sys.version.split()[0], + "gpu": cupy.cuda.runtime.getDeviceProperties(0)["name"].decode(), + "gpus": cupy.cuda.runtime.getDeviceCount(), + "host": socket.gethostname(), + "slurm_job": os.environ.get("SLURM_JOB_ID"), + } + + +# --------------------------------------------------------------------------- +# Experiment A: kernels +# --------------------------------------------------------------------------- + + +def _shapes( + kernel: str, b: int, bonds: tuple[int, ...], phys: int, eta: int +) -> tuple[tuple[int, ...], ...]: + """Operand shapes of one kernel call, as the planner sizes them.""" + cores = tuple((d, d, phys, phys) for d in bonds) + return { + "env": ((b, *bonds), (b, phys, phys), *cores), + "sketch": ((b, *bonds), *cores, (eta, *bonds)), + "project": ((b, eta, phys, phys), *cores, (eta, *bonds)), + }[kernel] + + +def _random(shape: tuple[int, ...]) -> Any: # noqa: ANN401 (a cupy array) + import cupy # noqa: PLC0415 + + scale = 1 / np.sqrt(2 * prod(shape[1:])) + out = cupy.random.standard_normal(shape) * scale + return out + 1j * (cupy.random.standard_normal(shape) * scale) + + +def _timed(fn: Callable[[], Any], warmups: int, runs: int) -> tuple[list[float], Any]: + """CUDA-event times of ``runs`` calls after ``warmups``, and the last output.""" + import cupy # noqa: PLC0415 + + out = None + for _ in range(warmups): + out = fn() + cupy.cuda.Device().synchronize() + times = [] + for _ in range(runs): + del out + start, end = cupy.cuda.Event(), cupy.cuda.Event() + start.record() + out = fn() + end.record() + end.synchronize() + times.append(cupy.cuda.get_elapsed_time(start, end) / 1e3) + return times, out + + +@app.command +def kernels( + bonds: tuple[int, ...] = (4, 4, 4000, 4), + phys: int = 4, + eta: int = 2000, + project: tuple[int, ...] = (8, 16, 32, 64, 128, 256), + sketch: tuple[int, ...] = (96, 192, 384), + env: tuple[int, ...] = (224,), + limit_gb: float = 50.0, + warmups: int = 2, + runs: int = 5, + output: Path = OUT / "kernels.jsonl", +) -> None: + """Experiment A: each kernel, batch and optimizer on random reference arrays. + + Both optimizers run in this process, each passed to `SiteKernels` and to + `peak_elements` as `src` passes it. A (kernel, optimizer, batch) whose planned + peak exceeds + ``limit_gb`` is skipped and recorded as such. + + Args: + bonds: The bond dimensions of the layers. + phys: The physical leg dimension. + eta: The sketch size (rows of the projected environment). + project: Projection batches (rows). + sketch: Sketch batches (columns). + env: Environment batches (columns). + limit_gb: Skip calls whose `peak_elements` exceeds this many GB. + warmups: Untimed calls before timing. + runs: Timed calls; the median is reported. + output: JSON lines file. + """ + import cupy # noqa: PLC0415 + + pool = cupy.get_default_memory_pool() + itemsize = np.dtype(np.complex128).itemsize + output.parent.mkdir(parents=True, exist_ok=True) + batches = {"env": env, "sketch": sketch, "project": project} + eqs = kernels_module.equations(len(bonds)) + with output.open("w") as f: + f.write(json.dumps({"meta": provenance()}) + "\n") + for kernel, sizes in batches.items(): + eq = {"env": eqs.ltr, "sketch": eqs.rtl_m, "project": eqs.rtl_s}[kernel] + cores = [_random((d, d, phys, phys)) for d in bonds] + proj = _random((eta, *bonds)) if kernel != "env" else None + for b in sizes: + shapes = _shapes(kernel, b, bonds, phys, eta) + lead = _random(shapes[0]) + omega = _random(shapes[1]) if kernel == "env" else None + outputs, args = {}, None + for opt in OPTIMIZERS: + start = perf_counter() + _, info = oe.contract_path(eq, *shapes, shapes=True, optimize=opt) + search = perf_counter() - start + start = perf_counter() + oe.contract_expression(eq, *shapes, optimize=opt) + build = perf_counter() - start + peak = kernels_module.peak_elements(eq, shapes, opt) + row: dict[str, Any] = { + "kernel": kernel, + "optimizer": opt, + "batch": b, + "flops": float(info.opt_cost), + "largest_intermediate": float(info.largest_intermediate), + "peak_elements": peak, + "peak_gb": peak * itemsize / GB, + "path": [list(p) for p in info.path], + "path_search_s": search, + "expression_build_s": build, + } + if peak * itemsize > limit_gb * GB: + row["status"] = "skipped: peak_elements above limit" + f.write(json.dumps(row) + "\n") + f.flush() + logger.info( + "%s b=%d %s: skipped (%.0f GB)", + kernel, + b, + opt, + row["peak_gb"], + ) + continue + sk = kernels_module.SiteKernels(len(bonds), opt) + args = { + "env": (lead, omega, cores), + "sketch": (lead, cores, proj), + "project": (lead, cores, proj), + }[kernel] + method = getattr(sk, kernel) + + def call(method=method, args=args) -> Any: # noqa: ANN401 + return method(*args) + + pool.free_all_blocks() + held = pool.used_bytes() + try: + times, out = _timed(call, warmups, runs) + except cupy.cuda.memory.OutOfMemoryError as err: + row["status"] = f"out of memory: {err}" + f.write(json.dumps(row) + "\n") + f.flush() + logger.info("%s b=%d %s: out of memory", kernel, b, opt) + pool.free_all_blocks() + continue + t = median(times) + row |= { + "status": "ok", + "times_s": times, + "median_s": t, + "per_unit_s": t / b, + "tflops": COMPLEX_FLOPS * row["flops"] / t / 1e12, + "pool_peak_gb": pool.total_bytes() / GB, + "pool_beyond_inputs_gb": (pool.total_bytes() - held) / GB, + "inputs_gb": held / GB, + } + outputs[opt] = out + del out + if len(outputs) == 2: + a, d = outputs["auto"], outputs["dp"] + row["rel_diff_vs_auto"] = float( + cupy.linalg.norm(a - d) / cupy.linalg.norm(a) + ) + f.write(json.dumps(row) + "\n") + f.flush() + logger.info( + "%s b=%d %s: %.4f s (%.3e s/unit), %.2f TFLOP/s, pool +%.1f GB%s", + kernel, + b, + opt, + t, + t / b, + row["tflops"], + row["pool_beyond_inputs_gb"], + f", diff {row['rel_diff_vs_auto']:.1e}" + if "rel_diff_vs_auto" in row + else "", + ) + del outputs, lead, omega, args + pool.free_all_blocks() + del cores, proj + pool.free_all_blocks() + + +# --------------------------------------------------------------------------- +# Experiments B and C: sweeps +# --------------------------------------------------------------------------- + + +class _Capture(logging.Handler): + """Collect the pass times and out-of-memory splits the sweep logs.""" + + def __init__(self) -> None: + super().__init__(logging.DEBUG) + self.records: list[str] = [] + + def emit(self, record: logging.LogRecord) -> None: + self.records.append(record.getMessage()) + + def passes(self, pattern: str) -> list[float]: + return [ + float(m.group(1)) + for msg in self.records + if (m := re.search(pattern + r": ([0-9.]+) s", msg)) + ] + + +@app.command +def sweep( + directory: Path, + *, + chi_out: int = 500, + gpu_memory: str = "60GB", + devices: int = 1, + repeats: int = 3, + seed: int = 0, + save_dir: Path | None = None, + optimize: str = "dp", + output: Path = OUT / "sweeps.jsonl", +) -> None: + """Experiments B and C: ``repeats`` sweeps with one path optimizer. + + Times the expression builds (path search included) separately: `_Contractions` + plans per `src` call, so every sweep pays them again. + + Args: + directory: A stack written by `bench_large` ``generate``. + chi_out: The output bond dimension. + gpu_memory: GPU budget of each device. + devices: The number of GPUs. + repeats: Sweeps to run. + seed: Seed of the sketch, the same for every sweep. + save_dir: Where to save the output of the last sweep, as ``.npz``. + optimize: The opt_einsum path optimizer, passed to `src`. + output: JSON lines file, appended to. + """ + import cupy # noqa: PLC0415 + + sys.path.insert(0, str(HERE)) + from bench_large import _load # noqa: PLC0415 + + opt = optimize + layers = _load(directory) + capture = _Capture() + logging.getLogger("src_method").addHandler(capture) + logging.getLogger("src_method").setLevel(logging.DEBUG) + + builds: list[float] = [] + real_build = kernels_module.contract_expression + + def timed_build(*args: Any, **kwargs: Any) -> Any: # noqa: ANN401 + start = perf_counter() + expr = real_build(*args, **kwargs) + builds.append(perf_counter() - start) + return expr + + kernels_module.contract_expression = timed_build + plans: list[Any] = [] + + def spy(*args: Any, **kwargs: Any) -> Any: # noqa: ANN401 + plans.append(make_plan(*args, **kwargs)) + return plans[-1] + + sweep_module.make_plan = spy + meta = provenance(opt) + output.parent.mkdir(parents=True, exist_ok=True) + out = None + with output.open("a") as f: + for rep in range(repeats): + for gpu in range(devices): + with cupy.cuda.Device(gpu): + cupy.get_default_memory_pool().free_all_blocks() + capture.records.clear() + builds.clear() + del out + start = perf_counter() + out = src( + *layers, + chi_out=chi_out, + dtype=layers[0][0].dtype, + seed=seed, + device="gpu", + resources=Resources(gpu_memory, None, None, devices=devices), + optimize=opt, + ) + seconds = perf_counter() - start + pools = [] + for gpu in range(devices): + with cupy.cuda.Device(gpu): + pools.append(cupy.get_default_memory_pool().total_bytes()) + plan = plans[-1] + tiers = [s.tier for s in plan.sites] + row = { + "experiment": "B" if devices == 1 else "C", + "optimizer": opt, + "devices": devices, + "repeat": rep, + "seconds": seconds, + "left_to_right_s": capture.passes("Left-to-right sweep"), + "right_to_left_s": capture.passes("Right-to-left sweep"), + "oom_splits": sum("Out of memory" in m for m in capture.records), + "expression_builds": len(builds), + "expression_build_s": sum(builds), + "batches": sorted( + {(s.env_batch, s.sketch_batch, s.project_batch) for s in plan.sites} + ), + "batches_per_site": [ + (s.env_batch, s.sketch_batch, s.project_batch) for s in plan.sites + ], + "planned_device_peak_gb": plan.device_peak / GB, + "pool_gb": [p / GB for p in pools], + "tiers": {t: tiers.count(t) for t in ("device", "host", "disk")}, + "chi_out": chi_out, + "gpu_memory": gpu_memory, + "stack": str(directory), + "meta": meta, + } + f.write(json.dumps(row) + "\n") + f.flush() + logger.info( + "%s rep %d: %.1f s, L2R %s, R2L %s, builds %d in %.2f s, batches %s, " + "pool %s GB, planned %.1f GB, oom splits %d", + opt, + rep, + seconds, + row["left_to_right_s"], + row["right_to_left_s"], + len(builds), + sum(builds), + row["batches"], + [round(p, 2) for p in row["pool_gb"]], + row["planned_device_peak_gb"], + row["oom_splits"], + ) + if save_dir is not None: + save_dir.mkdir(parents=True, exist_ok=True) + np.savez(save_dir / f"{opt}_g{devices}.npz", *out) + + +@app.command +def distance( + save_dir: Path, + devices: int = 1, + output: Path = OUT / "distances.jsonl", +) -> None: + """QR distance between the saved ``auto`` and ``dp`` outputs of `sweep`. + + Args: + save_dir: The ``save_dir`` of both sweeps. + devices: The GPU count of the sweeps to compare. + output: JSON lines file, appended to. + """ + sys.path.insert(0, str(HERE)) + from bench_large import _qr_distance # noqa: PLC0415 + + def load(opt: str) -> list[np.ndarray]: + with np.load(save_dir / f"{opt}_g{devices}.npz") as z: + return [z[f"arr_{i}"] for i in range(len(z.files))] + + d = _qr_distance(load("auto"), load("dp")) + logger.info("QR distance auto-dp on %d GPUs: %.3e", devices, d) + with output.open("a") as f: + f.write(json.dumps({"devices": devices, "qr_distance": d}) + "\n") + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s") + app() diff --git a/docs/content/docs/features/large-problems.mdx b/docs/content/docs/features/large-problems.mdx index 09f175a..11a6804 100644 --- a/docs/content/docs/features/large-problems.mdx +++ b/docs/content/docs/features/large-problems.mdx @@ -148,3 +148,24 @@ 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. With several devices, these lines start with the device they refer to. + +## Contraction paths + +Every kernel contracts its operands along the path that opt_einsum's `dp` +optimizer finds, the cheapest one, and the plan sizes batches for that same path. +The `optimize` argument of `src`, `apply` and `compress` takes any other path +optimizer opt_einsum accepts: a name such as `"auto"` or `"greedy"`, a name added +with `opt_einsum.paths.register_path_fn`, or an `opt_einsum.paths.PathOptimizer` +instance. An unknown name raises `ValueError` at once. + +```python notest +import opt_einsum as oe + +out = src(*stack, chi_out=2000, device="gpu", optimize="auto") +out = src(*stack, chi_out=2000, optimize=oe.DynamicProgramming(minimize="size")) +``` + +On an A100, `auto` took 1.9x as long as `dp` over a full sweep, and at large bonds +its projection path needs an intermediate too big for batches above 16 rows; see +the +[out-of-core benchmark](https://github.com/Algorithmiq/src-method/tree/main/benches/large). diff --git a/src/src_method/_kernels.py b/src/src_method/_kernels.py index 85da865..291670d 100644 --- a/src/src_method/_kernels.py +++ b/src/src_method/_kernels.py @@ -12,7 +12,7 @@ from functools import cache from itertools import count from math import prod -from typing import TYPE_CHECKING, NamedTuple +from typing import TYPE_CHECKING, NamedTuple, cast import opt_einsum as oe from opt_einsum import contract_expression, get_symbol @@ -21,11 +21,18 @@ from collections.abc import Sequence from opt_einsum.contract import ContractExpression + from opt_einsum.paths import PathOptimizer + from opt_einsum.typing import OptimizeKind from .utils import Array Shape = tuple[int, ...] +# The default path optimizer of the kernels. ``dp`` finds the cheapest path: on an +# A100 it halved the projection time against ``auto``, which at large bonds also +# plans an intermediate too big for any batch above 16 rows. +DEFAULT_OPTIMIZE = "dp" + class Equations(NamedTuple): """The einsum equations of one sweep, for a fixed stack depth.""" @@ -67,16 +74,22 @@ class _Contractions: 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. + + Args: + optimize: The opt_einsum path optimizer. """ - def __init__(self) -> None: + def __init__(self, optimize: str | PathOptimizer) -> None: + self._optimize = cast("OptimizeKind", optimize) self._compiled: dict[tuple[str, tuple[Shape, ...]], ContractExpression] = {} def __call__(self, eq: str, *operands: Array) -> Array: shapes = tuple(op.shape for op in operands) expr = self._compiled.get((eq, shapes)) if expr is None: - expr = self._compiled[eq, shapes] = contract_expression(eq, *shapes) + expr = self._compiled[eq, shapes] = contract_expression( + eq, *shapes, optimize=self._optimize + ) return expr(*operands) @@ -85,11 +98,14 @@ class SiteKernels: Args: depth: The number of layers of the stack. + optimize: The opt_einsum path optimizer, a name or a ``PathOptimizer``. """ - def __init__(self, depth: int) -> None: + def __init__( + self, depth: int, optimize: str | PathOptimizer = DEFAULT_OPTIMIZE + ) -> None: self.eqs = equations(depth) - self._contract = _Contractions() + self._contract = _Contractions(optimize) def env(self, env: Array, omega: Array, cores: Sequence[Array]) -> Array: """Advance a batch of sketch columns of the environment by one site. @@ -144,10 +160,15 @@ def first(self, cores: Sequence[Array], proj: Array) -> Array: @cache -def peak_elements(eq: str, shapes: tuple[Shape, ...]) -> int: +def peak_elements( + eq: str, + shapes: tuple[Shape, ...], + optimize: str | PathOptimizer = DEFAULT_OPTIMIZE, +) -> int: """Estimate the peak elements a contraction allocates beyond its inputs. - Walks the path `opt_einsum` picks for these shapes, the one `SiteKernels` runs. + Walks the path `opt_einsum` picks for these shapes with ``optimize``, the one + `SiteKernels` runs with it. 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 @@ -156,11 +177,14 @@ def peak_elements(eq: str, shapes: tuple[Shape, ...]) -> int: Args: eq: The einsum equation. shapes: The operand shapes. + optimize: The opt_einsum path optimizer. Returns: The peak number of elements. """ - _, info = oe.contract_path(eq, *shapes, shapes=True) + _, info = oe.contract_path( + eq, *shapes, shapes=True, optimize=cast("OptimizeKind", optimize) + ) sizes = info.size_dict # (elements, is_intermediate) for every operand still to be contracted. operands = [(prod(shape), False) for shape in shapes] diff --git a/src/src_method/_plan.py b/src/src_method/_plan.py index 0104e35..37e7834 100644 --- a/src/src_method/_plan.py +++ b/src/src_method/_plan.py @@ -22,13 +22,15 @@ import numpy as np -from ._kernels import equations, peak_elements +from ._kernels import DEFAULT_OPTIMIZE, equations, peak_elements from .utils import device_memory, host_memory_available, is_host, use_device if TYPE_CHECKING: import os from collections.abc import Callable + from opt_einsum.paths import PathOptimizer + from ._group import DeviceGroup Tier = Literal["device", "host", "disk"] @@ -294,6 +296,7 @@ def __init__( itemsize: int, n_sites: int, ranks: int = 1, + optimize: str | PathOptimizer = DEFAULT_OPTIMIZE, ) -> None: self.j, self.shapes, self.core_bytes, self.chi, self.e = ( j, @@ -302,7 +305,7 @@ def __init__( chi, itemsize, ) - self.ranks = ranks + self.ranks, self.optimize = ranks, optimize self.cols = ceil(chi / ranks) # The slice of the cores a device uploads before gathering the rest. self.chunk = ceil(core_bytes / ranks) if ranks > 1 else 0 @@ -325,7 +328,9 @@ def env(self, b: int, prefetch: int, *, staged_in: bool, staged_out: bool) -> in fixed = self._cores(prefetch) + self.cols * 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) + self.eqs.ltr, + ((b, *self.left), (b, self.up, self.down), *self.shapes), + self.optimize, ) return fixed + slices + peak * e @@ -342,6 +347,7 @@ def sketch(self, b: int, prefetch: int, *, staged: bool) -> int: peak_elements( self.eqs.rtl_m, ((b, *self.left), *self.shapes, (self.eta, *self.right)), + self.optimize, ) * e ) @@ -387,6 +393,7 @@ def project(self, b: int, prefetch: int) -> int: *self.shapes, (self.eta, *self.right), ), + self.optimize, ) * e ) @@ -398,7 +405,12 @@ def first(self, b: int, prefetch: int) -> int: fixed = ( self._cores(prefetch) + self.eta * self.b * e + self.a * chi * self.p * e ) - peak = peak_elements(self.eqs.first, (*self.shapes, (b, *self.right))) * e + peak = ( + peak_elements( + self.eqs.first, (*self.shapes, (b, *self.right)), self.optimize + ) + * e + ) return fixed + peak @@ -589,6 +601,7 @@ def make_plan( budgets: Budgets, *, devices: int = 1, + optimize: str | PathOptimizer = DEFAULT_OPTIMIZE, ) -> Plan: """Plan the batches and the environment tiers of one sweep. @@ -600,6 +613,8 @@ def make_plan( budgets: The resolved budgets. devices: The number of devices the sweep is split among; each runs the plan on its share of the sketch columns. + optimize: The opt_einsum path optimizer of the kernels, whose paths the + batches are sized for. Returns: The plan. Prefetching is dropped before giving up. @@ -619,6 +634,7 @@ def make_plan( itemsize=itemsize, n_sites=n_sites, ranks=devices, + optimize=optimize, ) for j, (shapes, core_bytes) in enumerate(zip(site_shapes, site_bytes)) ] diff --git a/src/src_method/_sweep.py b/src/src_method/_sweep.py index e9b0be4..f4b92de 100644 --- a/src/src_method/_sweep.py +++ b/src/src_method/_sweep.py @@ -36,7 +36,7 @@ from array_api_compat import array_namespace from ._group import DeviceGroup, blocks -from ._kernels import SiteKernels +from ._kernels import DEFAULT_OPTIMIZE, SiteKernels from ._plan import make_plan, resolve_budgets from ._sites import SiteSource, padded_shapes, site_bytes, sweep_order from ._store import EnvironmentStore @@ -56,6 +56,7 @@ from collections.abc import Callable, Iterator, Sequence from numpy.typing import DTypeLike + from opt_einsum.paths import PathOptimizer from ._plan import Budgets, Plan, Resources from ._tensor_train import Site, TrainKind @@ -168,6 +169,7 @@ def sweep( truncate_to: int | None = None, dtype: DTypeLike, resources: Resources | None = None, + optimize: str | PathOptimizer = DEFAULT_OPTIMIZE, ) -> list[np.ndarray]: """Contract and compress a stack in ket form with one SRC sweep. @@ -188,6 +190,7 @@ def sweep( dtype: The data type of the sketches. resources: Memory budgets, scratch space and devices; detected when ``None``. + optimize: The opt_einsum path optimizer of the kernels and of the plan. Returns: The site arrays of the compressed train as numpy arrays, right-canonical @@ -200,7 +203,13 @@ def sweep( group = DeviceGroup(device, resources.devices if resources else None) budgets = resolve_budgets(resources, group) plan = make_plan( - shapes, site_bytes(layers), chi_out, work, budgets, devices=group.size + shapes, + site_bytes(layers), + chi_out, + work, + budgets, + devices=group.size, + optimize=optimize, ) # Guarded: the per-site lists walk the whole plan, wasted work unless logged. if logger.isEnabledFor(logging.DEBUG): @@ -229,6 +238,7 @@ def sweep( sketches=_Sketches(prng, chi_out, dtype, group.size), work=work, depth=len(layers), + optimize=optimize, ) with SiteSource( layers, sweep_order(len(shapes)), depth=plan.prefetch, group=group @@ -316,10 +326,11 @@ def __init__( sketches: _Sketches, work: np.dtype, depth: int, + optimize: str | PathOptimizer, ) -> None: self.group, self.plan, self.budgets, self.shapes = group, plan, budgets, shapes self.chi_out, self.cutoff, self.sketches = chi_out, cutoff, sketches - self.work, self.depth = work, depth + self.work, self.depth, self.optimize = work, depth, optimize self.columns = blocks(chi_out, group.size) def _tag(self, rank: int) -> str: @@ -332,7 +343,7 @@ def run(self, rank: int, source: SiteSource) -> list[np.ndarray] | None: xp, tag = self.group.xp, self._tag(rank) lo, hi = self.columns[rank] env_shapes = [(hi - lo, *(s[0] for s in site)) for site in self.shapes] - kernels = _SplittingKernels(SiteKernels(self.depth), tag) + kernels = _SplittingKernels(SiteKernels(self.depth, self.optimize), tag) with ( device_pool_limit(xp, self.budgets.device_cap), EnvironmentStore( diff --git a/src/src_method/_validation.py b/src/src_method/_validation.py index 8e1aea1..cdabcd3 100644 --- a/src/src_method/_validation.py +++ b/src/src_method/_validation.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING import numpy as np +from opt_einsum.paths import PathOptimizer, get_path_fn if TYPE_CHECKING: from collections.abc import Sequence @@ -22,6 +23,7 @@ "validate_chi_out", "validate_cutoff", "validate_open_boundary", + "validate_optimize", ] @@ -93,6 +95,35 @@ def validate_cutoff(cutoff: float) -> None: raise ValueError(msg) +def validate_optimize(optimize: object) -> None: + """Reject anything opt_einsum cannot plan every kernel with. + + A name must resolve to an opt_einsum path optimizer, built in or added with + ``opt_einsum.paths.register_path_fn``. Explicit paths and booleans are refused: + the sweep contracts several equations, and no single path fits them all. + + Args: + optimize: The path optimizer. + + Raises: + TypeError: If `optimize` is neither a name nor a ``PathOptimizer``. + ValueError: If `optimize` names no opt_einsum path optimizer. + """ + if isinstance(optimize, PathOptimizer): + return + if not isinstance(optimize, str): + msg = ( + "optimize must be an opt_einsum path optimizer name or a " + f"PathOptimizer, got {optimize!r}." + ) + raise TypeError(msg) + try: + get_path_fn(optimize) + except KeyError as err: + msg = f"optimize names no opt_einsum path optimizer: {optimize!r}." + raise ValueError(msg) from err + + def validate_open_boundary( train: Sequence[Site], boundary_ndim: int, index: int ) -> None: diff --git a/src/src_method/apply.py b/src/src_method/apply.py index 1c4a4e4..5657331 100644 --- a/src/src_method/apply.py +++ b/src/src_method/apply.py @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING +from ._kernels import DEFAULT_OPTIMIZE from ._tensor_train import infer_kind from .stack import src @@ -18,6 +19,7 @@ import numpy as np from numpy.typing import DTypeLike + from opt_einsum.paths import PathOptimizer from ._plan import Resources from ._tensor_train import Site @@ -34,6 +36,7 @@ def apply( seed: int | None = None, device: str = "cpu", resources: Resources | None = None, + optimize: str | PathOptimizer = DEFAULT_OPTIMIZE, ) -> list[np.ndarray]: """Applies the Successive Randomized Compression (SRC) algorithm. @@ -62,6 +65,8 @@ def apply( 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`. + optimize: The opt_einsum path optimizer of the contractions, a name or a + ``PathOptimizer``; see `src_method.stack.src`. Defaults to ``"dp"``. Returns: The site arrays of the compressed tensor network (MPS or MPO), in @@ -70,13 +75,14 @@ def apply( Raises: TypeError: If ``chi_out`` is not an integer or the combination of input - tensor types is unsupported. + tensor types is unsupported, or ``optimize`` is neither a name nor a + ``PathOptimizer``. ValueError: If ``chi_out`` is not positive, if ``cutoff`` is not in ``[0.0, 1.0)``, if ``oversample`` is not finite and at least 1.0, if a train is not open-boundary, if the two trains differ in length or in the physical dimensions they join, if a - sub-three-site train is not exactly two sites, or if ``device`` is not - recognised. + sub-three-site train is not exactly two sites, if ``optimize`` names no + opt_einsum path optimizer, or if ``device`` is not recognised. ImportError: If ``device="gpu"`` but cupy is not installed. """ left_kind = infer_kind(left_tensor) @@ -98,4 +104,5 @@ def apply( seed=seed, device=device, resources=resources, + optimize=optimize, ) diff --git a/src/src_method/compress.py b/src/src_method/compress.py index 605c1f8..6d4e2d7 100644 --- a/src/src_method/compress.py +++ b/src/src_method/compress.py @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING +from ._kernels import DEFAULT_OPTIMIZE from .stack import src if TYPE_CHECKING: @@ -17,6 +18,7 @@ import numpy as np from numpy.typing import DTypeLike + from opt_einsum.paths import PathOptimizer from ._plan import Resources from ._tensor_train import Site @@ -32,6 +34,7 @@ def compress( seed: int | None = None, device: str = "cpu", resources: Resources | None = None, + optimize: str | PathOptimizer = DEFAULT_OPTIMIZE, ) -> list[np.ndarray]: """Applies the Successive Randomized Compression (SRC) algorithm. @@ -58,6 +61,8 @@ def compress( 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`. + optimize: The opt_einsum path optimizer of the contractions, a name or a + ``PathOptimizer``; see `src_method.stack.src`. Defaults to ``"dp"``. Returns: The site arrays of the compressed tensor network (MPS or MPO), in @@ -65,13 +70,13 @@ def compress( sites or more). Raises: - TypeError: If ``chi_out`` is not an integer or the input tensor type is - unsupported. + TypeError: If ``chi_out`` is not an integer, the input tensor type is + unsupported, or ``optimize`` is neither a name nor a ``PathOptimizer``. ValueError: If ``chi_out`` is not positive, if ``cutoff`` is not in ``[0.0, 1.0)``, if ``oversample`` is not finite and at least 1.0, if the train is not open-boundary, if a - sub-three-site train is not exactly two sites, or if ``device`` is not - recognised. + sub-three-site train is not exactly two sites, if ``optimize`` names no + opt_einsum path optimizer, or if ``device`` is not recognised. ImportError: If ``device="gpu"`` but cupy is not installed. """ return src( @@ -83,4 +88,5 @@ def compress( seed=seed, device=device, resources=resources, + optimize=optimize, ) diff --git a/src/src_method/stack.py b/src/src_method/stack.py index 5c15e95..212f98e 100644 --- a/src/src_method/stack.py +++ b/src/src_method/stack.py @@ -10,6 +10,7 @@ import logging from typing import TYPE_CHECKING +from ._kernels import DEFAULT_OPTIMIZE from ._sweep import sweep from ._tensor_train import ( MIN_SRC_SITES, @@ -17,7 +18,12 @@ exact_stack, normalize_stack, ) -from ._validation import sketch_size, validate_chi_out, validate_cutoff +from ._validation import ( + sketch_size, + validate_chi_out, + validate_cutoff, + validate_optimize, +) from .utils import default_rng, get_xp, sketch_dtype if TYPE_CHECKING: @@ -25,6 +31,7 @@ import numpy as np from numpy.typing import DTypeLike + from opt_einsum.paths import PathOptimizer from ._plan import Resources from ._tensor_train import Site @@ -46,6 +53,7 @@ def src( seed: int | None = None, device: str = "cpu", resources: Resources | None = None, + optimize: str | PathOptimizer = DEFAULT_OPTIMIZE, ) -> list[np.ndarray]: """Contract a stack of tensor trains and compress the result with SRC. @@ -98,6 +106,12 @@ def src( resources: Memory budgets and scratch space for the sweep (see `Resources`); every budget left unset is detected. Ignored for two-site stacks. + optimize: The opt_einsum path optimizer of the sweep's contractions: any + name opt_einsum knows (``"dp"``, ``"auto"``, ``"greedy"``, ... or one + added with ``opt_einsum.paths.register_path_fn``) or an + ``opt_einsum.paths.PathOptimizer`` instance. The memory plan sizes + batches for the same paths. Defaults to ``"dp"``, the cheapest paths. + Ignored for two-site stacks. Returns: The site arrays of the compressed train (MPS or MPO), as numpy arrays @@ -106,17 +120,20 @@ def src( Raises: TypeError: If ``chi_out`` is not an integer, if a train has an unrecognised - layout or an MPS sits anywhere other than at one end of the stack. + layout or an MPS sits anywhere other than at one end of the stack, + or if ``optimize`` is neither a name nor a ``PathOptimizer``. ValueError: If ``chi_out`` is not positive, if ``cutoff`` is not in ``[0.0, 1.0)``, if ``oversample`` is not finite and at least 1.0, if a train is not open-boundary, if the stack is empty, if the trains differ in length or in the physical dimensions they join, - if a sub-three-site stack is not exactly two sites, or if ``device`` is - not recognised. + if a sub-three-site stack is not exactly two sites, if ``optimize`` + names no opt_einsum path optimizer, or if ``device`` is not + recognised. ImportError: If ``device="gpu"`` but cupy is not installed. """ chi_out = validate_chi_out(chi_out) validate_cutoff(cutoff) + validate_optimize(optimize) chi_sketch = sketch_size(chi_out, oversample) get_xp(device) # Reject an unknown device, or a missing cupy, at once. prng = default_rng(seed) @@ -145,6 +162,7 @@ def src( truncate_to=chi_out if chi_sketch > chi_out else None, dtype=sketch_dtype(dtype, *layers), resources=resources, + optimize=optimize, ) logger.debug("SRC complete") return result diff --git a/tests/test_kernels.py b/tests/test_kernels.py index a3eb6fc..2162c16 100644 --- a/tests/test_kernels.py +++ b/tests/test_kernels.py @@ -1,6 +1,7 @@ """Test the batched site kernels and the peak-memory walk.""" import numpy as np +import opt_einsum as oe import pytest from src_method._kernels import SiteKernels, equations, peak_elements @@ -80,3 +81,21 @@ def test_peak_elements_of_a_matrix_chain(): shapes = ((2, 3), (3, 4), (4, 5)) assert peak_elements("ab,bc,cd->ad", shapes) == 46 + + +class FixedPath(oe.paths.PathOptimizer): + """Always return the same contraction path.""" + + def __init__(self, path): + self.path = path + + def __call__(self, inputs, output, size_dict, memory_limit=None): # noqa: ARG002 + return self.path + + +def test_peak_elements_follows_the_optimizer(): + # Contracting (3x4)(4x5) first holds its operands (12 + 20) and a 3x5 + # intermediate (15): 47, one more than the default path's 46. + shapes = ((2, 3), (3, 4), (4, 5)) + + assert peak_elements("ab,bc,cd->ad", shapes, FixedPath([(1, 2), (0, 1)])) == 47 diff --git a/tests/test_stack.py b/tests/test_stack.py index d38657d..0a6f52e 100644 --- a/tests/test_stack.py +++ b/tests/test_stack.py @@ -7,6 +7,7 @@ from functools import reduce import numpy as np +import opt_einsum as oe import pytest import src_method._sweep as sweep_module @@ -444,6 +445,58 @@ def test_src_validates_the_stack(rng): src(*make_stack("psi A psi", rng), chi_out=4) +# ------------------------- +# --- Contraction paths --- +# ------------------------- + + +@pytest.mark.parametrize( + "optimize", ["auto", "greedy", oe.DynamicProgramming(minimize="size")] +) +@pytest.mark.parametrize("spec", ["A B psi", "A B C", "phi A B"]) +def test_optimizer_does_not_change_the_result(spec, optimize, rng): + stack = make_stack(spec, rng) + + default = src(*stack, chi_out=4, seed=7) + other = src(*stack, chi_out=4, seed=7, optimize=optimize) + + assert rel_error(other, dense(default)) < 1e-10 + + +def test_registered_optimizer_is_used(rng, monkeypatch): + calls = [] + + def counting_greedy(*args, **kwargs): + calls.append(1) + return oe.paths.greedy(*args, **kwargs) + + # What opt_einsum.paths.register_path_fn does, undone after the test. + monkeypatch.setitem(oe.paths._PATH_OPTIONS, "src-test-greedy", counting_greedy) + stack = make_stack("A B psi", rng) + + out = src(*stack, chi_out=CHI_EXACT, optimize="src-test-greedy") + + assert calls + assert rel_error(out, dense_stack(*stack)) < 1e-10 + + +@pytest.mark.parametrize( + ("optimize", "error"), + [("fastest", ValueError), (True, TypeError), ([(0, 1), (0, 1)], TypeError)], +) +@pytest.mark.parametrize("entry", ["src", "apply", "compress"]) +def test_invalid_optimizer_raises(entry, optimize, error, rng): + A, psi = make_stack("A psi", rng) + call = { + "src": lambda: src(A, psi, chi_out=4, optimize=optimize), + "apply": lambda: apply(A, psi, 4, optimize=optimize), + "compress": lambda: compress(psi, 4, optimize=optimize), + }[entry] + + with pytest.raises(error, match="optimize"): + call() + + # ---------------- # --- Wrappers --- # ----------------