portc 0.6.0__py3-none-any.whl

portc/__init__.py

@@ -0,0 +1,29 @@
+"""Copyright (c) 2025 Algorithmiq Development Team. All rights reserved.
+
+portc: Parallel Out-of-Core Randomized Tensor-Network Compression.
+"""
+
+from __future__ import annotations
+
+import logging
+
+from ._plan import Resources
+from ._tensor_train import SiteLike
+from ._version import version as __version__
+from ._version import version_tuple as __version_tuple__
+from .apply import apply
+from .compress import compress
+from .stack import src
+
+# Stay silent unless the application configures logging.
+logging.getLogger(__name__).addHandler(logging.NullHandler())
+
+__all__ = [
+    "Resources",
+    "SiteLike",
+    "__version__",
+    "__version_tuple__",
+    "apply",
+    "compress",
+    "src",
+]

portc/_group.py

@@ -0,0 +1,295 @@
+"""The devices of one sweep, one thread per device, and the collectives between them.
+
+A sweep over several devices of a node runs in one process: `DeviceGroup.run`
+starts one thread per device, each running the same per-device function on its
+*rank* (its position in the group). Collectives are called by every rank in the
+same order. They exchange device arrays through peer copies, ordered with the
+kernels by events. A rank that fails aborts the group: every other rank waiting in
+a collective, or for an input core, raises `GroupAbortedError`, and `run` re-raises the
+original error.
+
+On the host backend the devices are simulated: the threads share the host, and the
+collectives are plain copies. The result is that of a single device; this mode
+exists to test the distributed logic without GPUs.
+"""
+
+from __future__ import annotations
+
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from typing import TYPE_CHECKING, Any, NamedTuple, TypeVar
+
+from ._plan import device_ids
+from .utils import (
+    copy_into,
+    current_device,
+    current_stream,
+    device_count,
+    enable_peer_access,
+    get_xp,
+    use_device,
+    use_stream,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Sequence
+
+    from .utils import Array
+
+T = TypeVar("T")
+
+
+class GroupAbortedError(RuntimeError):
+    """Raised on the ranks of a group after another rank failed."""
+
+
+def blocks(n: int, size: int) -> list[tuple[int, int]]:
+    """Split ``range(n)`` into ``size`` contiguous ``(lo, hi)`` blocks, in rank order.
+
+    Block sizes differ by at most one, the larger blocks first. Blocks may be empty
+    when ``n < size``.
+    """
+    base, extra = divmod(n, size)
+    out, lo = [], 0
+    for rank in range(size):
+        hi = lo + base + (rank < extra)
+        out.append((lo, hi))
+        lo = hi
+    return out
+
+
+class _Part(NamedTuple):
+    """What a rank publishes in a collective: an array and the event it is ready at."""
+
+    array: Any
+    ready: Any
+
+
+class DeviceGroup:
+    """The devices of one sweep and the collectives between them.
+
+    Args:
+        device: ``"cpu"`` or ``"gpu"``, as in `portc.stack.src`.
+        devices: A count ``n`` for the first ``n`` visible devices, or the device
+            ids in rank order; rank 0 is the root. Defaults to the current device.
+
+    Raises:
+        ValueError: If a device does not exist.
+    """
+
+    def __init__(self, device: str, devices: int | Sequence[int] | None = None) -> None:
+        self.xp = get_xp(device)
+        ids = device_ids(devices) or (current_device(self.xp),)
+        count = device_count(self.xp)
+        if count is not None and max(ids) >= count:
+            msg = f"Requested devices {list(ids)}, but only {count} are visible."
+            raise ValueError(msg)
+        self.devices = ids
+        self.size = len(self.devices)
+        # Uploading 1/size of each core and gathering the rest pays off only when
+        # the devices copy directly between them.
+        self.split_uploads = self.size > 1 and enable_peer_access(self.xp, ids)
+        self._barrier = threading.Barrier(self.size)
+        # Two sets of slots, alternating between consecutive collectives: a rank may
+        # publish into the next collective while slower ranks still read this one.
+        self._parts: list[list[Any]] = [[None] * self.size, [None] * self.size]
+        self._done: list[list[Any]] = [[None] * self.size, [None] * self.size]
+        self._calls = [0] * self.size
+        self._on_abort: list[Callable[[], None]] = []
+        self.aborted = threading.Event()
+
+    def on_abort(self, callback: Callable[[], None]) -> None:
+        """Call ``callback`` when a rank fails, to release threads outside the ranks."""
+        self._on_abort.append(callback)
+
+    def abort(self) -> None:
+        """Wake every rank waiting in a collective or on a registered resource."""
+        if self.aborted.is_set():
+            return
+        self.aborted.set()
+        self._barrier.abort()
+        for callback in self._on_abort:
+            callback()
+
+    def run(self, fn: Callable[[int], T]) -> list[T]:
+        """Run ``fn(rank)`` on every device, each in its own thread.
+
+        A group of one runs ``fn(0)`` in the calling thread.
+
+        Args:
+            fn: The per-device function.
+
+        Returns:
+            The results, in rank order.
+
+        Raises:
+            Exception: The first error raised by a rank, after every rank stopped.
+        """
+        if self.size == 1:
+            with use_device(self.xp, self.devices[0]), use_stream(self.xp):
+                return [fn(0)]
+        first: list[BaseException] = []
+        lock = threading.Lock()
+
+        def worker(rank: int) -> T:
+            try:
+                with use_device(self.xp, self.devices[rank]), use_stream(self.xp):
+                    return fn(rank)
+            except BaseException as err:
+                secondary = isinstance(
+                    err, GroupAbortedError | threading.BrokenBarrierError
+                )
+                with lock:
+                    if not first and not secondary:
+                        first.append(err)
+                self.abort()
+                raise
+
+        with ThreadPoolExecutor(self.size, thread_name_prefix="src-device") as pool:
+            futures = [pool.submit(worker, rank) for rank in range(self.size)]
+            results, errors = [], []
+            for future in futures:
+                try:
+                    results.append(future.result())
+                except BaseException as err:  # noqa: BLE001  (re-raised below)
+                    errors.append(err)
+        if first:
+            raise first[0]
+        if errors:
+            raise errors[0]
+        return results
+
+    # --- collectives ---
+
+    def _wait(self) -> None:
+        try:
+            self._barrier.wait()
+        except threading.BrokenBarrierError:
+            raise GroupAbortedError from None
+
+    def _publish(self, rank: int, array: Any) -> tuple[int, list[_Part]]:  # noqa: ANN401
+        """Publish ``array`` and return the parts of every rank."""
+        parity = self._calls[rank] % 2
+        self._calls[rank] += 1
+        self._parts[parity][rank] = _Part(array, current_stream(self.xp).record())
+        self._wait()
+        return parity, list(self._parts[parity])
+
+    def _settle(self, rank: int, parity: int) -> None:
+        """Wait, on this rank's stream, until every rank has read what it published.
+
+        Until then the arrays published by this rank must not be overwritten; the
+        memory pool reuses freed blocks in stream order, so making the stream wait
+        is enough.
+        """
+        stream = current_stream(self.xp)
+        self._done[parity][rank] = stream.record()
+        self._wait()
+        for event in self._done[parity]:
+            stream.wait_event(event)
+        self._parts[parity][rank] = None
+
+    def _pull(self, part: _Part, dst: Array) -> None:
+        stream = current_stream(self.xp)
+        stream.wait_event(part.ready)
+        copy_into(dst, part.array, stream)
+
+    def allgather(self, rank: int, local: Array) -> Array:
+        """Concatenate the blocks of every rank along the first axis, on every rank.
+
+        Args:
+            rank: The calling rank.
+            local: This rank's block, C-contiguous on its device; the blocks may
+                differ along the first axis only.
+
+        Returns:
+            The concatenation in rank order, on this rank's device.
+        """
+        if self.size == 1:
+            return local
+        parity, parts = self._publish(rank, local)
+        full = self._empty_like_parts(parts)
+        lo = 0
+        for part in parts:
+            hi = lo + part.array.shape[0]
+            self._pull(part, full[lo:hi])
+            lo = hi
+        self._settle(rank, parity)
+        return full
+
+    def gather(self, rank: int, local: Array, root: int = 0) -> Array | None:
+        """Concatenate the blocks of every rank along the first axis, on ``root``.
+
+        Args:
+            rank: The calling rank.
+            local: This rank's block, C-contiguous on its device.
+            root: The rank that receives the result.
+
+        Returns:
+            The concatenation on ``root``, ``None`` on the other ranks.
+        """
+        if self.size == 1:
+            return local
+        parity, parts = self._publish(rank, local)
+        full = None
+        if rank == root:
+            full = self._empty_like_parts(parts)
+            lo = 0
+            for part in parts:
+                hi = lo + part.array.shape[0]
+                self._pull(part, full[lo:hi])
+                lo = hi
+        self._settle(rank, parity)
+        return full
+
+    def scatter(self, rank: int, full: Array | None, root: int = 0) -> Array:
+        """Give every rank its block (see `blocks`) of the first axis of ``full``.
+
+        Args:
+            rank: The calling rank.
+            full: The array to split, C-contiguous on ``root``; ignored elsewhere.
+            root: The rank that holds ``full``.
+
+        Returns:
+            This rank's block, on its device.
+        """
+        if self.size == 1:
+            assert full is not None  # noqa: S101  (the root holds it)
+            return full
+        parity, parts = self._publish(rank, full if rank == root else None)
+        source = parts[root].array
+        lo, hi = blocks(source.shape[0], self.size)[rank]
+        local = self.xp.empty((hi - lo, *source.shape[1:]), dtype=source.dtype)
+        self._pull(_Part(source[lo:hi], parts[root].ready), local)
+        self._settle(rank, parity)
+        return local
+
+    def alltoall(self, rank: int, sends: Sequence[Array]) -> Array:
+        """Send ``sends[s]`` to rank ``s``; concatenate what every rank sent here.
+
+        Args:
+            rank: The calling rank.
+            sends: One C-contiguous block per rank, on this rank's device; the
+                blocks sent to one rank may differ along the first axis only.
+
+        Returns:
+            The blocks sent to this rank, concatenated along the first axis in rank
+            order, on this rank's device.
+        """
+        if self.size == 1:
+            return sends[0]
+        parity, parts = self._publish(rank, list(sends))
+        mine = [_Part(part.array[rank], part.ready) for part in parts]
+        full = self._empty_like_parts(mine)
+        lo = 0
+        for part in mine:
+            hi = lo + part.array.shape[0]
+            self._pull(part, full[lo:hi])
+            lo = hi
+        self._settle(rank, parity)
+        return full
+
+    def _empty_like_parts(self, parts: Sequence[_Part]) -> Array:
+        first = parts[0].array
+        rows = sum(part.array.shape[0] for part in parts)
+        return self.xp.empty((rows, *first.shape[1:]), dtype=first.dtype)

portc/_kernels.py

@@ -0,0 +1,199 @@
+"""Site kernels of the SRC sweep and the memory they need.
+
+Every kernel is one of the contractions of the sweep restricted to a batch: a slice
+of the sketch index for the environments and the sketch, a slice of the rows of
+the new projected environment otherwise. Batching never changes the result beyond
+rounding, because the sketch index and those rows are free indices of their
+contractions.
+"""
+
+from __future__ import annotations
+
+from functools import cache
+from itertools import count
+from math import prod
+from typing import TYPE_CHECKING, NamedTuple, cast
+
+import opt_einsum as oe
+from opt_einsum import contract_expression, get_symbol
+
+if TYPE_CHECKING:
+    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."""
+
+    ltr: str
+    rtl_m: str
+    rtl_s: str
+    first: str
+
+
+@cache
+def equations(depth: int) -> Equations:
+    """Generate the sweep equations for a stack of ``depth`` layers.
+
+    Layer ``i`` at a site carries ``(a_i, b_i, x_i, x_{i+1})``: left and right
+    bonds, then its upper and lower physical legs, so that consecutive layers share
+    ``x``. The output legs are ``x_0`` (up) and ``x_depth`` (down).
+    """
+    symbols = map(get_symbol, count())
+    sketch, eta_right, eta_left = next(symbols), next(symbols), next(symbols)
+    left = "".join(next(symbols) for _ in range(depth))
+    right = "".join(next(symbols) for _ in range(depth))
+    phys = [next(symbols) for _ in range(depth + 1)]
+    up, down = phys[0], phys[-1]
+    layers = ",".join(
+        f"{left[i]}{right[i]}{phys[i]}{phys[i + 1]}" for i in range(depth)
+    )
+    return Equations(
+        ltr=f"{sketch}{left},{sketch}{up}{down},{layers}->{sketch}{right}",
+        rtl_m=f"{sketch}{left},{layers},{eta_right}{right}->{eta_right}{up}{down}{sketch}",
+        rtl_s=f"{eta_left}{eta_right}{up}{down},{layers},{eta_right}{right}->{eta_left}{left}",
+        first=f"{layers},{eta_right}{right}->{left}{eta_right}{up}{down}",
+    )
+
+
+class _Contractions:
+    """Compiled contractions for one sweep, keyed on equation and operand shapes.
+
+    Uniform bulk sites share one entry, so the path is planned once rather than at
+    every site. Kept per call: jagged bonds and short last batches add entries that
+    are not worth keeping.
+
+    Args:
+        optimize: The opt_einsum path optimizer.
+    """
+
+    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, optimize=self._optimize
+            )
+        return expr(*operands)
+
+
+class SiteKernels:
+    """The four contractions of the sweep, each applied to one batch.
+
+    Args:
+        depth: The number of layers of the stack.
+        optimize: The opt_einsum path optimizer, a name or a ``PathOptimizer``.
+    """
+
+    def __init__(
+        self, depth: int, optimize: str | PathOptimizer = DEFAULT_OPTIMIZE
+    ) -> None:
+        self.eqs = equations(depth)
+        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.
+
+        Args:
+            env: Columns ``lo:hi`` of ``C_j``, shape ``(b, *left_bonds)``.
+            omega: The same columns of the site's Gaussian tensor, ``(b, up, down)``.
+            cores: The padded cores of the site.
+
+        Returns:
+            Columns ``lo:hi`` of ``C_{j+1}``, shape ``(b, *right_bonds)``.
+        """
+        return self._contract(self.eqs.ltr, env, omega, *cores)
+
+    def sketch(self, env: Array, cores: Sequence[Array], proj: Array) -> Array:
+        """Sketch a batch of columns of the site's running core.
+
+        Args:
+            env: Columns ``lo:hi`` of ``C_j``, shape ``(b, *left_bonds)``.
+            cores: The padded cores of the site.
+            proj: The projected environment ``S``, ``(eta, *right_bonds)``.
+
+        Returns:
+            Columns ``lo:hi`` of the sketch, shape ``(eta, up, down, b)``.
+        """
+        return self._contract(self.eqs.rtl_m, env, *cores, proj)
+
+    def project(self, eta: Array, cores: Sequence[Array], proj: Array) -> Array:
+        """Project a batch of rows of the new projected environment.
+
+        Args:
+            eta: Rows ``lo:hi`` of the conjugated output core, ``(b, eta, up, down)``.
+            cores: The padded cores of the site.
+            proj: The projected environment ``S``, ``(eta, *right_bonds)``.
+
+        Returns:
+            Rows ``lo:hi`` of the new ``S``, shape ``(b, *left_bonds)``.
+        """
+        return self._contract(self.eqs.rtl_s, eta, *cores, proj)
+
+    def first(self, cores: Sequence[Array], proj: Array) -> Array:
+        """Contract the first site with a batch of rows of ``S``.
+
+        Args:
+            cores: The padded cores of the first site.
+            proj: Rows ``lo:hi`` of ``S``, ``(b, *right_bonds)``.
+
+        Returns:
+            The output core for those rows, ``(*left_bonds, b, up, down)``.
+        """
+        return self._contract(self.eqs.first, *cores, proj)
+
+
+@cache
+def peak_elements(
+    eq: str,
+    shapes: tuple[Shape, ...],
+    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 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
+    are not.
+
+    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=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]
+    peak = 0
+    for step in info.contraction_list:
+        positions, einsum_str = step[0], step[2]
+        live = sum(n for n, is_tmp in operands if is_tmp)
+        popped = [operands.pop(i) for i in positions]
+        out = prod(sizes[c] for c in einsum_str.split("->")[1])
+        peak = max(peak, live + out + sum(n for n, _ in popped))
+        operands.append((out, True))
+    return peak