Large problems

A single SRC sweep keeps, besides the input and output trains, one sketched environment per site. For stacks with a large bond, such as N . V . M . U with a bond of several thousand in M, these environments and the intermediates of the contractions no longer fit on a GPU, and often not in host memory either. portc then plans the sweep to the memory at hand:

• the input cores are read one site at a time, so a train can live on disk;
• every contraction runs in batches of sketch columns (or rows), sized to the GPU budget;
• each environment is kept on the GPU, in host memory or in a scratch directory, the most recent ones on the fastest tier.

None of this changes the result beyond floating-point rounding: the random draws and the mathematics are those of a sweep that fits in memory.

BudgetsCopy Anchor Link

The budgets are set with Resources, passed to src, apply or compress:

from portc import Resources, src

out = src(
    N, V, M, U,
    chi_out=2000,
    dtype=np.complex128,
    device="gpu",
    resources=Resources(gpu_memory="36GB", scratch_dir="/local/scratch"),
)

Every field left unset is detected when the call starts:

Sizes are byte counts or strings: "36GB" is 36 * 10**9 bytes, "36GiB" is 36 * 2**30. On the CPU there is a single budget, host_memory. Host memory is measured when the call starts, so inputs already held in memory are not counted twice.

The sweep is planned to fit the GPU budget, counting the bytes its arrays use. CuPy's memory pool also holds blocks that are split and only partly in use, so during the sweep the pool is capped at the GPU budget plus the margin of max(10%, 1 GiB) of the device. That margin is the one detection leaves free, so with a detected budget the cap is the free device memory; the previous limit is restored afterwards. When the pool fragments beyond the margin, which large batches can do, a batch that runs out of memory is computed in halves, and that kernel keeps the finer split for the rest of the sweep; the result does not change. If a single site does not fit the budget even with batches of one column, src raises MemoryError naming the site.

Inputs on diskCopy Anchor Link

A train is any sequence of per-site array-likes with shape, dtype, ndim and np.asarray support (the SiteLike protocol, exported for type annotations): NumPy arrays, np.memmap, zarr or HDF5 datasets. Each core is read when the sweep reaches its site, once per pass, and a background thread reads the next site while the current one is computed. Raw .npy files opened with np.load(path, mmap_mode="r") are the fastest option; compressed formats may be limited by decompression.

The scratch directoryCopy Anchor Link

Environments that fit in neither budget are written to a per-process directory, <scratch_dir>/src-<pid>-<id>/, one file per site. Use node-local disk: at the reference size of 50 sites, D_M = 4000 and chi_out = 2000 in complex128, up to 410 GB are written and read back once. The directory is removed when the call returns, also after an error; a process killed with SIGKILL leaves it behind, and its name identifies the process.

Several GPUsCopy Anchor Link

On a node with several GPUs, one sweep can be split among them:

from portc import Resources, src

out = src(
    N, V, M, U,
    chi_out=2000,
    dtype=np.complex128,
    device="gpu",
    resources=Resources(devices=4, scratch_dir="/local/scratch"),
)

devices is a count, for the first n visible GPUs, or a list of device ids, the first one leading. Everything runs in the calling process, one thread per GPU, so nothing changes for the caller: the call returns the output cores as with one GPU, errors are raised as usual, and the inputs are read once and shared. Each thread runs on a non-blocking CUDA stream of its own GPU, not on CuPy's legacy default stream, where a copy between two GPUs waits for the copy in the other direction and the exchanges between GPUs would get half the bandwidth. The stream first waits for the work already queued on the caller's current stream, so inputs given as CuPy arrays may still be in flight.

The sketch index is split: each GPU owns a contiguous block of the chi_out sketch columns and computes only those columns of every environment and of every sketch, so the environments, the dominant memory, are split among the GPUs too. At every site of the right-to-left pass, the sketch is orthonormalized where it lies when chi_out is at least 1024: the GPUs swap their blocks of columns for blocks of rows, and a shifted CholeskyQR (at most three passes, each exchanging one chi_out x chi_out Gram matrix per GPU) builds their rows of the orthonormal basis without moving the sketch again. Below that size, and at the sites whose sketch is too ill-conditioned for CholeskyQR (numerically rank-deficient, as at small bonds or early in an evolution), the first GPU gathers the sketch and runs a Householder QR instead. Either way, each GPU then has a block of rows of the output core, projects its rows of the next projected environment, and those rows are gathered on every GPU. The random draws are those of one GPU, so the result is that of one GPU up to rounding and a gauge: the two QRs span the same basis but may pick a different orthonormal frame of it, which changes the cores, not the operator they represent. With DEBUG logging, each sweep reports at how many sites it tried CholeskyQR and how many of those fell back.

What is not split:

• every GPU holds a whole site: the cores, the projected environment and the output core of the site, so the largest site that fits is that of one GPU;
• every core is read from the source once, but every GPU needs a copy. When the GPUs reach each other directly (NVLink, or PCIe peer-to-peer), each uploads a slice of the core and gathers the rest from the others; otherwise each uploads the whole core;
• below chi_out = 1024, and where CholeskyQR falls back, the QR runs on the first GPU while the others wait.

These costs grow more slowly with the problem than the contractions do, so large problems, the ones this is for, scale close to linearly, while small ones may run slower on several GPUs than on one.

The budgets apply as follows: gpu_memory is the budget of each GPU (detected: the smallest free memory of the GPUs), while host_memory and scratch_dir are shared by all of them. Every GPU writes its environments to its own scratch directory.

On the CPU, devices simulates the GPUs with threads that share the host; the result is the same, and the mode exists to test the multi-GPU sweep without GPUs.

Reading the planCopy Anchor Link

With DEBUG logging on for the portc logger (see Logging), every sweep logs its plan as SRC plan: ...:

• prefetch: whether the next site is read ahead;
• devices: the device ids the sweep is split among;
• device peak, host peak, disk: the planned peaks, in bytes; the device peak is that of each device, the others are totals;
• scratch: the scratch directory, when anything spills to disk;
• tiers: where the environment of each site lives (device, host or disk);
• batches: for each site, the batch of the environment, sketch and projection steps.

It then logs the time of each pass and, at the end, SRC stalls: the seconds the sweep waited for input cores and for environments. If they are a large fraction of the run, the disk is too slow for the compute. With several devices, these lines start with the device they refer to.

Contraction pathsCopy Anchor Link

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.

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.