From 259dcd2999a36d2529d39871f58b360c89976778 Mon Sep 17 00:00:00 2001 From: Panadestein Date: Tue, 27 Oct 2026 11:12:11 +0100 Subject: [PATCH 1/5] =?UTF-8?q?feat(kernels):=20=E2=9A=A1=20contract=20alo?= =?UTF-8?q?ng=20opt=5Feinsum's=20dp=20paths=20by=20default?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SiteKernels now plans every contraction with opt_einsum's "dp" optimizer instead of the default "auto" heuristic, and peak_elements walks the same path, so the planner sizes batches for the path the kernel executes. On an A100, "auto" made the projection kernel cost twice dp's flops and, at the reference bonds, plan an 8.2e9-element intermediate that held the projection batch to 16 rows. With dp the full D_M = 1000 sweep runs 1.95x faster on one GPU and 1.86x on four, with the same batches and outputs equal to rounding. SRC_EINSUM_OPTIMIZE selects another opt_einsum optimizer (e.g. "auto") at import; an unknown name raises ValueError. Documented under large problems. Assisted-by: Claude Code:claude-opus-5-5 --- docs/content/docs/features/large-problems.mdx | 11 +++++ src/src_method/_kernels.py | 48 +++++++++++++++++-- tests/test_kernels.py | 21 +++++++- 3 files changed, 75 insertions(+), 5 deletions(-) diff --git a/docs/content/docs/features/large-problems.mdx b/docs/content/docs/features/large-problems.mdx index 09f175a..5da1e5d 100644 --- a/docs/content/docs/features/large-problems.mdx +++ b/docs/content/docs/features/large-problems.mdx @@ -148,3 +148,14 @@ 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 cheapest path that opt_einsum's +`dp` optimizer finds, and the plan sizes batches for that same path. The +environment variable `SRC_EINSUM_OPTIMIZE` selects another opt_einsum optimizer, +such as `auto`. It is read once, when `src_method` is imported, and an unknown +name raises `ValueError` at import. On an A100, `auto` took 1.9x as long 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..89c967a 100644 --- a/src/src_method/_kernels.py +++ b/src/src_method/_kernels.py @@ -9,10 +9,11 @@ from __future__ import annotations +import os from functools import cache from itertools import count from math import prod -from typing import TYPE_CHECKING, NamedTuple +from typing import TYPE_CHECKING, Literal, NamedTuple, cast, get_args import opt_einsum as oe from opt_einsum import contract_expression, get_symbol @@ -26,6 +27,42 @@ Shape = tuple[int, ...] +# The named path optimizers of opt_einsum. +PathOptimizer = Literal[ + "auto", + "auto-hq", + "dp", + "greedy", + "optimal", + "branch-all", + "branch-2", + "random-greedy", + "random-greedy-128", +] + + +def _path_optimizer() -> PathOptimizer: + """Read ``SRC_EINSUM_OPTIMIZE``, ``"dp"`` when unset. + + Raises: + ValueError: If it names no opt_einsum path optimizer. + """ + name = os.environ.get("SRC_EINSUM_OPTIMIZE", "dp") + if name not in get_args(PathOptimizer): + msg = ( + f"SRC_EINSUM_OPTIMIZE must be one of {get_args(PathOptimizer)}, " + f"got {name!r}." + ) + raise ValueError(msg) + return cast("PathOptimizer", name) + + +# The path optimizer of every kernel, read once at import. ``dp`` finds the +# cheapest path: on an A100 it halves the projection time against ``auto``, which +# at large bonds also plans an intermediate too big for any batch above 16 rows. +# ``SRC_EINSUM_OPTIMIZE`` overrides it, e.g. ``auto`` to compare. +OPTIMIZE: PathOptimizer = _path_optimizer() + class Equations(NamedTuple): """The einsum equations of one sweep, for a fixed stack depth.""" @@ -76,7 +113,9 @@ 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=OPTIMIZE + ) return expr(*operands) @@ -147,7 +186,8 @@ def first(self, cores: Sequence[Array], proj: Array) -> Array: 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. + Walks the path `opt_einsum` picks for these shapes with `OPTIMIZE`, 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 @@ -160,7 +200,7 @@ def peak_elements(eq: str, shapes: tuple[Shape, ...]) -> int: Returns: The peak number of elements. """ - _, info = oe.contract_path(eq, *shapes, shapes=True) + _, info = oe.contract_path(eq, *shapes, shapes=True, optimize=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/tests/test_kernels.py b/tests/test_kernels.py index a3eb6fc..c27cc69 100644 --- a/tests/test_kernels.py +++ b/tests/test_kernels.py @@ -3,7 +3,7 @@ import numpy as np import pytest -from src_method._kernels import SiteKernels, equations, peak_elements +from src_method._kernels import SiteKernels, _path_optimizer, equations, peak_elements def cores(rng, dtype, *, left=(3, 2), right=(4, 3), up=2, mid=3, down=2): @@ -80,3 +80,22 @@ def test_peak_elements_of_a_matrix_chain(): shapes = ((2, 3), (3, 4), (4, 5)) assert peak_elements("ab,bc,cd->ad", shapes) == 46 + + +def test_path_optimizer_defaults_to_dp(monkeypatch): + monkeypatch.delenv("SRC_EINSUM_OPTIMIZE", raising=False) + + assert _path_optimizer() == "dp" + + +def test_path_optimizer_reads_the_environment(monkeypatch): + monkeypatch.setenv("SRC_EINSUM_OPTIMIZE", "auto") + + assert _path_optimizer() == "auto" + + +def test_path_optimizer_rejects_unknown_names(monkeypatch): + monkeypatch.setenv("SRC_EINSUM_OPTIMIZE", "fastest") + + with pytest.raises(ValueError, match="SRC_EINSUM_OPTIMIZE"): + _path_optimizer() From 5efda5380c7da6cb6c99945b0f1d3913ff4b72c5 Mon Sep 17 00:00:00 2001 From: Panadestein Date: Tue, 27 Oct 2026 11:12:11 +0100 Subject: [PATCH 2/5] =?UTF-8?q?chore(bench):=20=F0=9F=93=88=20measure=20au?= =?UTF-8?q?to=20against=20dp=20einsum=20paths=20on=20Leonardo?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add benches/large/bench_einsum_dp.py: kernel timings at the reference shapes (both optimizers in one process, CUDA events, planned peaks), full sweeps with the optimizer of the process, and the QR distance between their outputs. The results of the A100-SXM-64GB runs are in benches/large/README.md. Assisted-by: Claude Code:claude-sonnet-5-5 --- benches/large/README.md | 49 +++ benches/large/bench_einsum_dp.py | 667 +++++++++++++++++++++++++++++++ 2 files changed, 716 insertions(+) create mode 100644 benches/large/bench_einsum_dp.py diff --git a/benches/large/README.md b/benches/large/README.md index f2a4802..5c912a9 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` + +`SRC_EINSUM_OPTIMIZE` selects the opt_einsum path optimizer of the kernels, and +the planner sizes batches for the same path. It is read at import and 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..995a277 --- /dev/null +++ b/benches/large/bench_einsum_dp.py @@ -0,0 +1,667 @@ +"""opt_einsum ``auto`` against ``dp`` paths for the SRC kernels, on the GPU. + +Four 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 + the optimizer that ``SRC_EINSUM_OPTIMIZE`` selects, 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; +- ``report``: tables and plots from the JSON lines. +""" + +from __future__ import annotations + +import json +import logging +import os +import re +import socket +import subprocess +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 = HERE / "leonardo" / "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() -> dict[str, Any]: + """Commit, versions, job and GPU of this process.""" + import cupy # noqa: PLC0415 (GPU only) + + def git(*args: str) -> str: + return subprocess.run( # noqa: S603 + ["git", *args], # noqa: S607 + cwd=HERE, + capture_output=True, + text=True, + check=False, + ).stdout.strip() + + return { + "commit": git("rev-parse", "HEAD"), + "src_dirty": bool(git("status", "--porcelain", "--", "src")), + "optimize_env": os.environ.get("SRC_EINSUM_OPTIMIZE"), + "optimize_module": kernels_module.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: `OPTIMIZE` is switched between them and + the `peak_elements` cache cleared, so each runs the path it would select from + ``SRC_EINSUM_OPTIMIZE``. 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: + kernels_module.OPTIMIZE = opt + kernels_module.peak_elements.cache_clear() + 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) + 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)) + 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() + kernels_module.OPTIMIZE = kernels_module._path_optimizer() # noqa: SLF001 + + +# --------------------------------------------------------------------------- +# 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, + output: Path = OUT / "sweeps.jsonl", +) -> None: + """Experiments B and C: ``repeats`` sweeps with the optimizer of the process. + + 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``. + output: JSON lines file, appended to. + """ + import cupy # noqa: PLC0415 + + sys.path.insert(0, str(HERE)) + from bench_large import _load # noqa: PLC0415 + + opt = kernels_module.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() + 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), + ) + 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") + + +# --------------------------------------------------------------------------- +# Report +# --------------------------------------------------------------------------- + +COLORS = {"auto": "#2a78d6", "dp": "#eb6834"} +MARKERS = {"auto": "o", "dp": "s"} + + +def _jsonl(path: Path) -> list[dict[str, Any]]: + return ( + [json.loads(x) for x in path.read_text().splitlines()] if path.exists() else [] + ) + + +@app.command +def report(folder: Path = OUT) -> None: + """Write ``summary.md`` and the per-kernel plot from the JSON lines. + + A hand-written ``conclusion.md`` in ``folder``, if present, is appended. + + Args: + folder: Folder with ``kernels.jsonl``, ``sweeps.jsonl`` and ``distances.jsonl``. + """ + import matplotlib as mpl # noqa: PLC0415 + + mpl.use("Agg") + import matplotlib.pyplot as plt # noqa: PLC0415 + + krows = _jsonl(folder / "kernels.jsonl") + kmeta = next((r["meta"] for r in krows if "meta" in r), None) + krows = [r for r in krows if "meta" not in r] + srows = _jsonl(folder / "sweeps.jsonl") + drows = _jsonl(folder / "distances.jsonl") + + lines = ["# opt_einsum `auto` against `dp` for the SRC kernels", ""] + lines += [ + "## Provenance", + "", + "| experiment | job | commit | cupy | numpy | opt_einsum | gpu |", + "|---|---|---|---|---|---|---|", + ] + metas = ([("A", kmeta)] if kmeta else []) + [ + (f"{r['experiment']} {r['optimizer']}", r["meta"]) + for r in srows + if r["repeat"] == 0 + ] + for name, m in metas: + lines.append( + f"| {name} | {m['slurm_job']} | {m['commit'][:10]}" + f"{' (dirty src)' if m['src_dirty'] else ''} | {m['cupy']} | {m['numpy']}" + f" | {m['opt_einsum']} | {m['gpus']}x {m['gpu']} |" + ) + lines.append("") + + if krows: + lines += [ + "## A: kernels at the reference shapes", + "", + "D = (4, 4, 4000, 4), physical legs 4, eta = 2000, complex128, one A100. " + "Median of 5 runs after 2 warm-ups, CUDA events. `flops` is opt_einsum's " + "count; TFLOP/s counts 4 real flops per opt_einsum op (complex128). " + "`pool +GB` is the pool growth beyond the inputs. `search` is " + "`contract_path` alone, paid once per shape and sweep.", + "", + "| kernel | batch | optimizer | flops | largest interm. | planned peak GB" + " | median s | s / unit | TFLOP/s | pool +GB | search ms | rel. diff |", + "|---|---|---|---|---|---|---|---|---|---|---|---|", + ] + for r in krows: + ok = r["status"] == "ok" + lines.append( + f"| {r['kernel']} | {r['batch']} | {r['optimizer']} | {r['flops']:.2e}" + f" | {r['largest_intermediate']:.2e} | {r['peak_gb']:.1f}" + + ( + f" | {r['median_s']:.4f} | {r['per_unit_s']:.2e} | {r['tflops']:.2f}" + f" | {r['pool_beyond_inputs_gb']:.1f}" + if ok + else f" | {r['status']} | | | " + ) + + f" | {1e3 * r['path_search_s']:.1f}" + + ( + f" | {r['rel_diff_vs_auto']:.1e} |" + if "rel_diff_vs_auto" in r + else " | |" + ) + ) + lines.append("") + fig, axes = plt.subplots(1, 3, figsize=(12, 3.8), layout="constrained") + for ax, kernel in zip(axes, ("project", "sketch", "env"), strict=True): + notes, top = [], 0.0 + for opt in OPTIMIZERS: + pts = [ + r + for r in krows + if r["kernel"] == kernel + and r["optimizer"] == opt + and r["status"] == "ok" + ] + top = max([top, *(r["per_unit_s"] * 1e3 for r in pts)]) + # auto is drawn larger, underneath: the two often coincide. + ax.plot( + [r["batch"] for r in pts], + [r["per_unit_s"] * 1e3 for r in pts], + color=COLORS[opt], + marker=MARKERS[opt], + markersize=10 if opt == "auto" else 6, + markeredgecolor="white", + linewidth=2, + label=opt, + zorder=2 if opt == "auto" else 3, + ) + skipped = [ + r["batch"] + for r in krows + if r["kernel"] == kernel + and r["optimizer"] == opt + and r["status"] != "ok" + ] + if skipped: + notes.append(f"{opt} skipped: {', '.join(map(str, skipped))}") + if notes: + ax.text( + 0.98, + 0.97, + "\n".join([*notes, "(planned peak > 50 GB)"]), + transform=ax.transAxes, + ha="right", + va="top", + fontsize=8, + color="#52514e", + ) + batches = sorted({r["batch"] for r in krows if r["kernel"] == kernel}) + unit = "row" if kernel == "project" else "column" + ax.set_xscale("log", base=2) + ax.set_xticks(batches, [str(x) for x in batches]) + ax.minorticks_off() + ax.set_xlim(batches[0] / 1.5, batches[-1] * 1.5) + ax.set_xlabel(f"batch ({unit}s)") + ax.set_ylabel(f"ms per {unit}") + ax.set_title(kernel, fontsize=10) + ax.set_ylim(0, 1.45 * top) + ax.grid(True, color="#e6e5e0", linewidth=0.8) + for side in ("top", "right"): + ax.spines[side].set_visible(False) + ax.legend(frameon=False, loc="lower left") + fig.savefig(folder / "kernels_time_per_unit.png", dpi=150) + plt.close(fig) + lines += ["![kernels](kernels_time_per_unit.png)", ""] + + for exp, title in ( + ("B", "B: full sweep, one A100"), + ("C", "C: full sweep, 4 A100s"), + ): + rows = [r for r in srows if r["experiment"] == exp] + if not rows: + continue + lines += [ + f"## {title}", + "", + "D_M = 1000, chi_out = 500, 50 sites, complex128, 60 GB GPU budget per " + "device. Per repeat: wall time, the passes (max over devices), expression " + "builds (path search included), out-of-memory splits.", + "", + "| optimizer | repeat | wall s | left-to-right s | right-to-left s | builds (s)" + " | batches (env, sketch, project) | pool GB | planned peak GB | tiers | OOM splits |", + "|---|---|---|---|---|---|---|---|---|---|---|", + ] + for r in rows: + lines.append( + f"| {r['optimizer']} | {r['repeat']} | {r['seconds']:.1f}" + f" | {max(r['left_to_right_s'], default=float('nan')):.1f}" + f" | {max(r['right_to_left_s'], default=float('nan')):.1f}" + f" | {r['expression_builds']} ({r['expression_build_s']:.2f})" + f" | {', '.join(str(tuple(b)) for b in r['batches'])}" + f" | {', '.join(f'{p:.1f}' for p in r['pool_gb'])}" + f" | {r['planned_device_peak_gb']:.1f}" + f" | {r['tiers']['device']}/{r['tiers']['host']}/{r['tiers']['disk']}" + f" | {r['oom_splits']} |" + ) + lines += [ + "", + "| optimizer | median s | range s | median L2R s | median R2L s |", + "|---|---|---|---|---|", + ] + for opt in OPTIMIZERS: + sub = [r for r in rows if r["optimizer"] == opt] + if not sub: + continue + secs = [r["seconds"] for r in sub] + lines.append( + f"| {opt} | {median(secs):.1f} | {min(secs):.1f}-{max(secs):.1f}" + f" | {median(max(r['left_to_right_s']) for r in sub):.1f}" + f" | {median(max(r['right_to_left_s']) for r in sub):.1f} |" + ) + devices = rows[0]["devices"] + for d in drows: + if d["devices"] == devices: + lines += [ + "", + f"QR distance between the `auto` and `dp` outputs: " + f"{d['qr_distance']:.1e}.", + ] + lines.append("") + + conclusion = folder / "conclusion.md" + if conclusion.exists(): + lines += [conclusion.read_text().rstrip(), ""] + (folder / "summary.md").write_text("\n".join(lines)) + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s") + app() From 0c6086d265cc71542c231e193aaba1ddabd23fcf Mon Sep 17 00:00:00 2001 From: Panadestein Date: Tue, 27 Oct 2026 11:37:58 +0100 Subject: [PATCH 3/5] =?UTF-8?q?chore(bench):=20=F0=9F=94=A5=20drop=20the?= =?UTF-8?q?=20report=20and=20plots=20from=20the=20einsum=20benchmark?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The results live in benches/large/README.md, like the other benches, so the report command, its plotting and the default results folder inside the repository go. The commands write JSON lines under ./einsum_dp instead. Assisted-by: Claude Code:claude-sonnet-5-5 --- benches/large/bench_einsum_dp.py | 219 +------------------------------ 1 file changed, 5 insertions(+), 214 deletions(-) diff --git a/benches/large/bench_einsum_dp.py b/benches/large/bench_einsum_dp.py index 995a277..3d941dd 100644 --- a/benches/large/bench_einsum_dp.py +++ b/benches/large/bench_einsum_dp.py @@ -1,6 +1,6 @@ """opt_einsum ``auto`` against ``dp`` paths for the SRC kernels, on the GPU. -Four commands: +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 @@ -8,8 +8,9 @@ - ``sweep``: run the SRC of a ``generate``-d stack of `bench_large` a few times with the optimizer that ``SRC_EINSUM_OPTIMIZE`` selects, 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; -- ``report``: tables and plots from the JSON lines. +- ``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 @@ -45,7 +46,7 @@ app = cyclopts.App(help="Measure opt_einsum auto against dp for the SRC kernels.") HERE = Path(__file__).parent -OUT = HERE / "leonardo" / "einsum_dp" +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 @@ -452,216 +453,6 @@ def load(opt: str) -> list[np.ndarray]: f.write(json.dumps({"devices": devices, "qr_distance": d}) + "\n") -# --------------------------------------------------------------------------- -# Report -# --------------------------------------------------------------------------- - -COLORS = {"auto": "#2a78d6", "dp": "#eb6834"} -MARKERS = {"auto": "o", "dp": "s"} - - -def _jsonl(path: Path) -> list[dict[str, Any]]: - return ( - [json.loads(x) for x in path.read_text().splitlines()] if path.exists() else [] - ) - - -@app.command -def report(folder: Path = OUT) -> None: - """Write ``summary.md`` and the per-kernel plot from the JSON lines. - - A hand-written ``conclusion.md`` in ``folder``, if present, is appended. - - Args: - folder: Folder with ``kernels.jsonl``, ``sweeps.jsonl`` and ``distances.jsonl``. - """ - import matplotlib as mpl # noqa: PLC0415 - - mpl.use("Agg") - import matplotlib.pyplot as plt # noqa: PLC0415 - - krows = _jsonl(folder / "kernels.jsonl") - kmeta = next((r["meta"] for r in krows if "meta" in r), None) - krows = [r for r in krows if "meta" not in r] - srows = _jsonl(folder / "sweeps.jsonl") - drows = _jsonl(folder / "distances.jsonl") - - lines = ["# opt_einsum `auto` against `dp` for the SRC kernels", ""] - lines += [ - "## Provenance", - "", - "| experiment | job | commit | cupy | numpy | opt_einsum | gpu |", - "|---|---|---|---|---|---|---|", - ] - metas = ([("A", kmeta)] if kmeta else []) + [ - (f"{r['experiment']} {r['optimizer']}", r["meta"]) - for r in srows - if r["repeat"] == 0 - ] - for name, m in metas: - lines.append( - f"| {name} | {m['slurm_job']} | {m['commit'][:10]}" - f"{' (dirty src)' if m['src_dirty'] else ''} | {m['cupy']} | {m['numpy']}" - f" | {m['opt_einsum']} | {m['gpus']}x {m['gpu']} |" - ) - lines.append("") - - if krows: - lines += [ - "## A: kernels at the reference shapes", - "", - "D = (4, 4, 4000, 4), physical legs 4, eta = 2000, complex128, one A100. " - "Median of 5 runs after 2 warm-ups, CUDA events. `flops` is opt_einsum's " - "count; TFLOP/s counts 4 real flops per opt_einsum op (complex128). " - "`pool +GB` is the pool growth beyond the inputs. `search` is " - "`contract_path` alone, paid once per shape and sweep.", - "", - "| kernel | batch | optimizer | flops | largest interm. | planned peak GB" - " | median s | s / unit | TFLOP/s | pool +GB | search ms | rel. diff |", - "|---|---|---|---|---|---|---|---|---|---|---|---|", - ] - for r in krows: - ok = r["status"] == "ok" - lines.append( - f"| {r['kernel']} | {r['batch']} | {r['optimizer']} | {r['flops']:.2e}" - f" | {r['largest_intermediate']:.2e} | {r['peak_gb']:.1f}" - + ( - f" | {r['median_s']:.4f} | {r['per_unit_s']:.2e} | {r['tflops']:.2f}" - f" | {r['pool_beyond_inputs_gb']:.1f}" - if ok - else f" | {r['status']} | | | " - ) - + f" | {1e3 * r['path_search_s']:.1f}" - + ( - f" | {r['rel_diff_vs_auto']:.1e} |" - if "rel_diff_vs_auto" in r - else " | |" - ) - ) - lines.append("") - fig, axes = plt.subplots(1, 3, figsize=(12, 3.8), layout="constrained") - for ax, kernel in zip(axes, ("project", "sketch", "env"), strict=True): - notes, top = [], 0.0 - for opt in OPTIMIZERS: - pts = [ - r - for r in krows - if r["kernel"] == kernel - and r["optimizer"] == opt - and r["status"] == "ok" - ] - top = max([top, *(r["per_unit_s"] * 1e3 for r in pts)]) - # auto is drawn larger, underneath: the two often coincide. - ax.plot( - [r["batch"] for r in pts], - [r["per_unit_s"] * 1e3 for r in pts], - color=COLORS[opt], - marker=MARKERS[opt], - markersize=10 if opt == "auto" else 6, - markeredgecolor="white", - linewidth=2, - label=opt, - zorder=2 if opt == "auto" else 3, - ) - skipped = [ - r["batch"] - for r in krows - if r["kernel"] == kernel - and r["optimizer"] == opt - and r["status"] != "ok" - ] - if skipped: - notes.append(f"{opt} skipped: {', '.join(map(str, skipped))}") - if notes: - ax.text( - 0.98, - 0.97, - "\n".join([*notes, "(planned peak > 50 GB)"]), - transform=ax.transAxes, - ha="right", - va="top", - fontsize=8, - color="#52514e", - ) - batches = sorted({r["batch"] for r in krows if r["kernel"] == kernel}) - unit = "row" if kernel == "project" else "column" - ax.set_xscale("log", base=2) - ax.set_xticks(batches, [str(x) for x in batches]) - ax.minorticks_off() - ax.set_xlim(batches[0] / 1.5, batches[-1] * 1.5) - ax.set_xlabel(f"batch ({unit}s)") - ax.set_ylabel(f"ms per {unit}") - ax.set_title(kernel, fontsize=10) - ax.set_ylim(0, 1.45 * top) - ax.grid(True, color="#e6e5e0", linewidth=0.8) - for side in ("top", "right"): - ax.spines[side].set_visible(False) - ax.legend(frameon=False, loc="lower left") - fig.savefig(folder / "kernels_time_per_unit.png", dpi=150) - plt.close(fig) - lines += ["![kernels](kernels_time_per_unit.png)", ""] - - for exp, title in ( - ("B", "B: full sweep, one A100"), - ("C", "C: full sweep, 4 A100s"), - ): - rows = [r for r in srows if r["experiment"] == exp] - if not rows: - continue - lines += [ - f"## {title}", - "", - "D_M = 1000, chi_out = 500, 50 sites, complex128, 60 GB GPU budget per " - "device. Per repeat: wall time, the passes (max over devices), expression " - "builds (path search included), out-of-memory splits.", - "", - "| optimizer | repeat | wall s | left-to-right s | right-to-left s | builds (s)" - " | batches (env, sketch, project) | pool GB | planned peak GB | tiers | OOM splits |", - "|---|---|---|---|---|---|---|---|---|---|---|", - ] - for r in rows: - lines.append( - f"| {r['optimizer']} | {r['repeat']} | {r['seconds']:.1f}" - f" | {max(r['left_to_right_s'], default=float('nan')):.1f}" - f" | {max(r['right_to_left_s'], default=float('nan')):.1f}" - f" | {r['expression_builds']} ({r['expression_build_s']:.2f})" - f" | {', '.join(str(tuple(b)) for b in r['batches'])}" - f" | {', '.join(f'{p:.1f}' for p in r['pool_gb'])}" - f" | {r['planned_device_peak_gb']:.1f}" - f" | {r['tiers']['device']}/{r['tiers']['host']}/{r['tiers']['disk']}" - f" | {r['oom_splits']} |" - ) - lines += [ - "", - "| optimizer | median s | range s | median L2R s | median R2L s |", - "|---|---|---|---|---|", - ] - for opt in OPTIMIZERS: - sub = [r for r in rows if r["optimizer"] == opt] - if not sub: - continue - secs = [r["seconds"] for r in sub] - lines.append( - f"| {opt} | {median(secs):.1f} | {min(secs):.1f}-{max(secs):.1f}" - f" | {median(max(r['left_to_right_s']) for r in sub):.1f}" - f" | {median(max(r['right_to_left_s']) for r in sub):.1f} |" - ) - devices = rows[0]["devices"] - for d in drows: - if d["devices"] == devices: - lines += [ - "", - f"QR distance between the `auto` and `dp` outputs: " - f"{d['qr_distance']:.1e}.", - ] - lines.append("") - - conclusion = folder / "conclusion.md" - if conclusion.exists(): - lines += [conclusion.read_text().rstrip(), ""] - (folder / "summary.md").write_text("\n".join(lines)) - - if __name__ == "__main__": logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s") app() From df71cad4213471c65a39aecf1bed383bd1488c90 Mon Sep 17 00:00:00 2001 From: Panadestein Date: Tue, 27 Oct 2026 11:58:41 +0100 Subject: [PATCH 4/5] =?UTF-8?q?chore(bench):=20=E2=99=BB=EF=B8=8F=20take?= =?UTF-8?q?=20the=20commit=20from=20the=20hatch-vcs=20version?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The version of src_method already records the commit and whether the tree was dirty, so the einsum benchmark no longer shells out to git. Assisted-by: Claude Code:claude-sonnet-5-5 --- benches/large/bench_einsum_dp.py | 14 +------------- 1 file changed, 1 insertion(+), 13 deletions(-) diff --git a/benches/large/bench_einsum_dp.py b/benches/large/bench_einsum_dp.py index 3d941dd..a35ee46 100644 --- a/benches/large/bench_einsum_dp.py +++ b/benches/large/bench_einsum_dp.py @@ -20,7 +20,6 @@ import os import re import socket -import subprocess import sys from importlib.metadata import version from math import prod @@ -55,21 +54,10 @@ def provenance() -> dict[str, Any]: - """Commit, versions, job and GPU of this process.""" + """Versions (``src_method`` carries the commit), job and GPU of this process.""" import cupy # noqa: PLC0415 (GPU only) - def git(*args: str) -> str: - return subprocess.run( # noqa: S603 - ["git", *args], # noqa: S607 - cwd=HERE, - capture_output=True, - text=True, - check=False, - ).stdout.strip() - return { - "commit": git("rev-parse", "HEAD"), - "src_dirty": bool(git("status", "--porcelain", "--", "src")), "optimize_env": os.environ.get("SRC_EINSUM_OPTIMIZE"), "optimize_module": kernels_module.OPTIMIZE, "src_method": src_method.__version__, From 507b7683b12fc9a525375ee58995e028b9dfc042 Mon Sep 17 00:00:00 2001 From: Panadestein Date: Tue, 27 Oct 2026 12:12:25 +0100 Subject: [PATCH 5/5] =?UTF-8?q?feat(kernels):=20=E2=9C=A8=20select=20the?= =?UTF-8?q?=20einsum=20path=20optimizer=20through=20the=20API?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the SRC_EINSUM_OPTIMIZE environment variable with an optimize argument on src, apply and compress, defaulting to "dp". It takes any name opt_einsum resolves, including those added with register_path_fn, or a PathOptimizer instance, so new optimizers work without a list to maintain. Validation defers to opt_einsum's own resolver. The value reaches the kernels and the memory plan, so batches are still sized for the path that runs. Assisted-by: Claude Code:claude-opus-5-5 --- README.md | 4 +- benches/large/README.md | 6 +- benches/large/bench_einsum_dp.py | 31 ++++---- docs/content/docs/features/large-problems.mdx | 24 ++++-- src/src_method/_kernels.py | 74 ++++++++----------- src/src_method/_plan.py | 24 +++++- src/src_method/_sweep.py | 19 ++++- src/src_method/_validation.py | 31 ++++++++ src/src_method/apply.py | 13 +++- src/src_method/compress.py | 14 +++- src/src_method/stack.py | 26 ++++++- tests/test_kernels.py | 26 +++---- tests/test_stack.py | 53 +++++++++++++ 13 files changed, 241 insertions(+), 104 deletions(-) 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 5c912a9..b2e8198 100644 --- a/benches/large/README.md +++ b/benches/large/README.md @@ -119,9 +119,9 @@ Speed-up and efficiency are against the single-GPU run: ### Einsum paths: `auto` against `dp` -`SRC_EINSUM_OPTIMIZE` selects the opt_einsum path optimizer of the kernels, and -the planner sizes batches for the same path. It is read at import and defaults -to `dp` since these measurements; the results above ran with `auto`. +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 diff --git a/benches/large/bench_einsum_dp.py b/benches/large/bench_einsum_dp.py index a35ee46..c5c521c 100644 --- a/benches/large/bench_einsum_dp.py +++ b/benches/large/bench_einsum_dp.py @@ -6,8 +6,8 @@ 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 - the optimizer that ``SRC_EINSUM_OPTIMIZE`` selects, logging the plan, the pass - times and the device peak, and save the last output; + 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``. @@ -53,13 +53,12 @@ COMPLEX_FLOPS = 4 -def provenance() -> dict[str, Any]: +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_env": os.environ.get("SRC_EINSUM_OPTIMIZE"), - "optimize_module": kernels_module.OPTIMIZE, + "optimize": optimize, "src_method": src_method.__version__, "numpy": np.__version__, "cupy": cupy.__version__, @@ -132,9 +131,9 @@ def kernels( ) -> None: """Experiment A: each kernel, batch and optimizer on random reference arrays. - Both optimizers run in this process: `OPTIMIZE` is switched between them and - the `peak_elements` cache cleared, so each runs the path it would select from - ``SRC_EINSUM_OPTIMIZE``. A (kernel, optimizer, batch) whose planned peak exceeds + 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: @@ -168,15 +167,13 @@ def kernels( omega = _random(shapes[1]) if kernel == "env" else None outputs, args = {}, None for opt in OPTIMIZERS: - kernels_module.OPTIMIZE = opt - kernels_module.peak_elements.cache_clear() 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) + peak = kernels_module.peak_elements(eq, shapes, opt) row: dict[str, Any] = { "kernel": kernel, "optimizer": opt, @@ -201,7 +198,7 @@ def kernels( row["peak_gb"], ) continue - sk = kernels_module.SiteKernels(len(bonds)) + sk = kernels_module.SiteKernels(len(bonds), opt) args = { "env": (lead, omega, cores), "sketch": (lead, cores, proj), @@ -260,7 +257,6 @@ def call(method=method, args=args) -> Any: # noqa: ANN401 pool.free_all_blocks() del cores, proj pool.free_all_blocks() - kernels_module.OPTIMIZE = kernels_module._path_optimizer() # noqa: SLF001 # --------------------------------------------------------------------------- @@ -296,9 +292,10 @@ def sweep( 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 the optimizer of the process. + """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. @@ -311,6 +308,7 @@ def sweep( 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 @@ -318,7 +316,7 @@ def sweep( sys.path.insert(0, str(HERE)) from bench_large import _load # noqa: PLC0415 - opt = kernels_module.OPTIMIZE + opt = optimize layers = _load(directory) capture = _Capture() logging.getLogger("src_method").addHandler(capture) @@ -341,7 +339,7 @@ def spy(*args: Any, **kwargs: Any) -> Any: # noqa: ANN401 return plans[-1] sweep_module.make_plan = spy - meta = provenance() + meta = provenance(opt) output.parent.mkdir(parents=True, exist_ok=True) out = None with output.open("a") as f: @@ -360,6 +358,7 @@ def spy(*args: Any, **kwargs: Any) -> Any: # noqa: ANN401 seed=seed, device="gpu", resources=Resources(gpu_memory, None, None, devices=devices), + optimize=opt, ) seconds = perf_counter() - start pools = [] diff --git a/docs/content/docs/features/large-problems.mdx b/docs/content/docs/features/large-problems.mdx index 5da1e5d..11a6804 100644 --- a/docs/content/docs/features/large-problems.mdx +++ b/docs/content/docs/features/large-problems.mdx @@ -151,11 +151,21 @@ these lines start with the device they refer to. ## Contraction paths -Every kernel contracts its operands along the cheapest path that opt_einsum's -`dp` optimizer finds, and the plan sizes batches for that same path. The -environment variable `SRC_EINSUM_OPTIMIZE` selects another opt_einsum optimizer, -such as `auto`. It is read once, when `src_method` is imported, and an unknown -name raises `ValueError` at import. On an A100, `auto` took 1.9x as long over a -full sweep, and at large bonds its projection path needs an intermediate too big -for batches above 16 rows; see the +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 89c967a..291670d 100644 --- a/src/src_method/_kernels.py +++ b/src/src_method/_kernels.py @@ -9,11 +9,10 @@ from __future__ import annotations -import os from functools import cache from itertools import count from math import prod -from typing import TYPE_CHECKING, Literal, NamedTuple, cast, get_args +from typing import TYPE_CHECKING, NamedTuple, cast import opt_einsum as oe from opt_einsum import contract_expression, get_symbol @@ -22,46 +21,17 @@ 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 named path optimizers of opt_einsum. -PathOptimizer = Literal[ - "auto", - "auto-hq", - "dp", - "greedy", - "optimal", - "branch-all", - "branch-2", - "random-greedy", - "random-greedy-128", -] - - -def _path_optimizer() -> PathOptimizer: - """Read ``SRC_EINSUM_OPTIMIZE``, ``"dp"`` when unset. - - Raises: - ValueError: If it names no opt_einsum path optimizer. - """ - name = os.environ.get("SRC_EINSUM_OPTIMIZE", "dp") - if name not in get_args(PathOptimizer): - msg = ( - f"SRC_EINSUM_OPTIMIZE must be one of {get_args(PathOptimizer)}, " - f"got {name!r}." - ) - raise ValueError(msg) - return cast("PathOptimizer", name) - - -# The path optimizer of every kernel, read once at import. ``dp`` finds the -# cheapest path: on an A100 it halves the projection time against ``auto``, which -# at large bonds also plans an intermediate too big for any batch above 16 rows. -# ``SRC_EINSUM_OPTIMIZE`` overrides it, e.g. ``auto`` to compare. -OPTIMIZE: PathOptimizer = _path_optimizer() +# 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): @@ -104,9 +74,13 @@ 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: @@ -114,7 +88,7 @@ def __call__(self, eq: str, *operands: Array) -> Array: expr = self._compiled.get((eq, shapes)) if expr is None: expr = self._compiled[eq, shapes] = contract_expression( - eq, *shapes, optimize=OPTIMIZE + eq, *shapes, optimize=self._optimize ) return expr(*operands) @@ -124,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. @@ -183,11 +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 with `OPTIMIZE`, 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 @@ -196,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, optimize=OPTIMIZE) + _, 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 c27cc69..2162c16 100644 --- a/tests/test_kernels.py +++ b/tests/test_kernels.py @@ -1,9 +1,10 @@ """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, _path_optimizer, equations, peak_elements +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): @@ -82,20 +83,19 @@ def test_peak_elements_of_a_matrix_chain(): assert peak_elements("ab,bc,cd->ad", shapes) == 46 -def test_path_optimizer_defaults_to_dp(monkeypatch): - monkeypatch.delenv("SRC_EINSUM_OPTIMIZE", raising=False) +class FixedPath(oe.paths.PathOptimizer): + """Always return the same contraction path.""" - assert _path_optimizer() == "dp" + def __init__(self, path): + self.path = path + def __call__(self, inputs, output, size_dict, memory_limit=None): # noqa: ARG002 + return self.path -def test_path_optimizer_reads_the_environment(monkeypatch): - monkeypatch.setenv("SRC_EINSUM_OPTIMIZE", "auto") - assert _path_optimizer() == "auto" - - -def test_path_optimizer_rejects_unknown_names(monkeypatch): - monkeypatch.setenv("SRC_EINSUM_OPTIMIZE", "fastest") +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)) - with pytest.raises(ValueError, match="SRC_EINSUM_OPTIMIZE"): - _path_optimizer() + 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 --- # ----------------