diffct 2.1.0__tar.gz → 2.2.0__tar.gz

{diffct-2.1.0 → diffct-2.2.0}/CHANGELOG.md

@@ -7,6 +7,49 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [2.2.0] - 2026-10-11
+
+### Added
+
+- Operation-specific automatic schedules and slab/block comparisons for
+  native projection and matched backprojection, with bounded CUDA pilots
+  and per-call execution/copy metadata.
+- Detector-pixel tiling and explicit `ParameterizedSurface` batches using
+  global pixel/view coordinates, saved parameter/frame values and local VJP
+  reduction. Legacy two-argument surface callbacks remain compatible.
+- Caller-owned numerical `project_into`/`backproject_into` outputs, tensor
+  block access and mapped `.npy` stores for data exceeding host RAM.
+- Distributed `partition="space"` with fixed volume slabs, ownership-aware
+  autograd/norms and disk CGLS state/checkpoints without full-volume gathering.
+- Independent CUDA tests and bounded formal protocol/ownership verification.
+
+### Changed
+
+- GPU backprojection accumulators now span all contributing view batches,
+  with final transfer after tile completion. Contiguous native volume layouts
+  and ray buffers are reused across batches.
+- Transfer execution uses bounded pinned slots, separate streams and events;
+  all extra staging/cached GPU buffers are included in capacity estimates.
+  Local GPU and NCCL reductions occur before completed-result downloads.
+- Example CGLS norms accumulate FP64 blocks and vector updates reuse buffers.
+  The disk example keeps all three volume states plus residual/projection state
+  blockwise and retains `q` until the global `alpha` decision.
+
+### Fixed
+
+- Tensor CGLS ranks agree on CPU fallback before jointly replanning unequal
+  spatial shards, preserving collective order.
+- Parameterized detector surfaces use a common arithmetic dtype of at least
+  float32, supporting mixed frame dtypes and finite half-precision inputs.
+- CUDA store uploads wait for newly generated blocks on their source stream,
+  including transfers within one GPU.
+- Single-slot execution retires prior native computation before allocating
+  the next tile, keeping deferred CUDA allocations within its admission bound.
+- Mapped input and output stores reject hard links to the same file before writes.
+- Example PSNR supports CPU reconstruction results with CUDA reference volumes.
+- Deferred CPU output writes retain their submission order across spatial tiles.
+- Coronal slice displays preserve the measured detector row orientation.
+
 ## [2.1.0] - 2026-10-10
 
 ### Added

{diffct-2.1.0 → diffct-2.2.0}/CITATION.cff

@@ -13,8 +13,8 @@ authors:
     given-names: Yipeng
     email: yipeng.sun@fau.de
     alias: sypsyp97
-version: 2.1.0
-date-released: "2026-10-10"
+version: 2.2.0
+date-released: "2026-10-11"
 license: Apache-2.0
 repository-code: "https://github.com/sypsyp97/diffct"
 url: "https://sypsyp97.github.io/diffct/"

{diffct-2.1.0 → diffct-2.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: diffct
-Version: 2.1.0
+Version: 2.2.0
 Summary: Differentiable CUDA projectors for CT: arbitrary trajectories, multi-GPU and multi-node execution, geometry gradients
 Project-URL: Homepage, https://github.com/sypsyp97/diffct
 Project-URL: Documentation, https://sypsyp97.github.io/diffct/
@@ -65,8 +65,6 @@ Description-Content-Type: text/markdown
   <a href="#citation">Citation</a>
 </p>
 
-> **Install** with the CUDA extra, `pip install "diffct[cu12]"` or `pip install "diffct[cu13]"`, after PyTorch. A plain `pip install diffct` needs a system CUDA Toolkit.
->
 > **Version 2.0** replaces the circular-orbit API of 1.x with per-view trajectories. Code written for 1.x needs changes; see the [migration notes](https://github.com/sypsyp97/diffct/blob/main/docs/MIGRATION.md).
 > For Apple Silicon, use the [MLX port by Linda-Sophie Schneider](https://github.com/Linda-SophieSchneider/DiffCT-MLX).
 
@@ -75,10 +73,10 @@ Description-Content-Type: text/markdown
 To our knowledge, diffct is the only open-source GPU CT library that combines arbitrary per-view trajectories, autograd through the volume, first-order gradients through the acquisition geometry, and built-in multi-GPU and multi-node execution. We compared it with LEAP, TIGRE, ASTRA/tomosipo, DiffDRR and others in October 2026.
 
 - **Arbitrary trajectories.** Each view has its own source, detector centre and detector axes. Circular, helical, saddle, sinusoidal, random and calibrated scans use the same code. Trajectory tensors with `requires_grad=True` receive geometry gradients for calibration.
-- **Curved detectors.** A `detector_surface(u, v)` callback defines each pixel's position for arc, cylindrical or other parameterized surfaces. Projection and matched backprojection support first-order surface gradients. See [the surface API](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#parameterized-detector-surfaces).
+- **Curved detectors.** `ParameterizedSurface` defines native arc, cylindrical or arbitrary pixels from global coordinates and explicit parameters, sampled in bounded view/pixel batches. Legacy `detector_surface(u, v)` callbacks remain available. Projection and matched backprojection support first-order surface gradients. See [the surface API](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#parameterized-detector-surfaces).
 - **Matched operators.** `project()` and `backproject()` form an exact adjoint pair for the cell-constant Siddon model. Both support PyTorch autograd, with volume/sinogram gradients and Hessian-vector products.
-- **Large volumes.** CPU volumes and sinograms automatically stream spatial tiles and view batches through CUDA. Full arrays and data gradients stay on CPU; tile sizes follow the available device memory and can be overridden. See [chunked execution](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#chunked-execution).
-- **Multi-GPU and multi-node.** Use `devices=[0, 1, 2, 3]` in one process, or one process per GPU with torchrun and NCCL. Views are partitioned. CPU-backed volumes stay in host RAM while GPUs compute tiles; distributed ranks retain replicated-volume semantics. Speedup depends on workload and communication costs.
+- **Large volumes.** Automatic spatial, view and pixel blocks with bounded transfer buffers control working memory. `project_into`/`backproject_into` accept caller-owned tensor or disk stores; tensor autograd keeps complete data on its input device. See [chunked execution](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#chunked-execution).
+- **Multi-GPU and multi-node.** Use `devices=[0, 1, 2, 3]` in one process, or torchrun and NCCL across ranks. Default `partition="views"` shards projections and replicates volumes. `partition="space"` gives each rank an owned volume slab and combines ray contributions. Speedup depends on workload and communication costs.
 - **Analytical helpers.** `diffct.analytical` provides ramp filters (ram-lak, shepp-logan, cosine, hamming, hann), fan, cone and Parker weights, and FBP and FDK backprojection.
 
 Capabilities, limits and isocenter rules: [docs/REFERENCE.md](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#capabilities-and-limits).
@@ -140,7 +138,7 @@ loss.backward()                     # x.grad = A^T (A x - y)
 
 **Geometry gradients.** Set `requires_grad=True` on trajectory tensors before you build `Projector`. See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#geometry-gradients).
 
-**Several GPUs and nodes.** Use `devices=[0, 1, 2, 3]` in one process, or one process per GPU with torchrun for one or more nodes. See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#several-gpus-and-nodes) and [DISTRIBUTED.md](https://github.com/sypsyp97/diffct/blob/main/docs/DISTRIBUTED.md).
+**Several GPUs and nodes.** See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#several-gpus-and-nodes) and [DISTRIBUTED.md](https://github.com/sypsyp97/diffct/blob/main/docs/DISTRIBUTED.md).
 
 ```python
 A = Projector(trajectory, (128, 128, 128), (384, 256), detector_spacing=0.8, devices=[0, 1, 2, 3])
@@ -161,29 +159,36 @@ sinogram.square().mean().backward()          # CPU volume.grad
 Run `python examples/chunked_reconstruction.py` for CPU-backed CGLS and measured
 CUDA allocation. Chunking preserves the Siddon model, with float32 summation
 differences; smaller blocks can increase transfer and launch overhead. Full
-arrays and iterative state still need enough host RAM.
+tensor arrays and iterative state still need enough host RAM. For disk-backed
+state and sharded checkpoints, run
+`python examples/disk_reconstruction.py --output disk-run`; for fixed rank-local
+volume slabs use `--partition space` under torchrun. Run
+`python examples/benchmark_execution.py` to compare schedules and copy payloads.
 
 ## Performance
 
 <p align="center">
   <picture>
-    <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_dark.png">
-    <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_light.png" width="100%" alt="Time per forward, adjoint and CGLS iteration on 1, 4 and 8 GPUs, 64^3 and 128^3">
+    <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_dark.png?v=2.2.0">
+    <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_light.png?v=2.2.0" width="100%" alt="Time per forward, adjoint and CGLS iteration on 1, 4 and 8 GPUs, 512^3 volume">
   </picture>
 </p>
 
-| 128³ volume | 1 GPU | 4 GPUs, 1 node | 8 GPUs, 2 nodes |
+| 512³ volume | 1 GPU | 4 GPUs, 1 node | 8 GPUs, 2 nodes |
 |---|---:|---:|---:|
-| Forward projection | 8.45 ms | 2.66 ms (3.2×) | 1.68 ms (5.0×) |
-| Adjoint (backprojection) | 20.66 ms | 5.89 ms (3.5×) | 4.62 ms (4.5×) |
-| CGLS iteration | 30.08 ms | 8.08 ms (3.7×) | 5.59 ms (5.4×) |
+| Forward projection | 493.05 ms | 132.22 ms (3.7×) | 73.18 ms (6.7×) |
+| Adjoint (backprojection) | 1400.80 ms | 369.29 ms (3.8×) | 265.87 ms (5.3×) |
+| CGLS iteration | 2054.36 ms | 547.06 ms (3.8×) | 373.38 ms (5.5×) |
 
-Circular trajectory, 360 views, detector (2n, 1.5n) cells at pitch 1.25, A100 64 GB GPUs. Small volumes scale less: at 64³, one CGLS iteration takes 4.63 ms on 1 GPU, 1.57 ms on 4 and 1.56 ms on 8. Raw data: [docs/assets/scaling.json](https://github.com/sypsyp97/diffct/blob/main/docs/assets/scaling.json).
+diffct 2.2.0, FP32, A100 64 GB GPUs; circular cone beam, 360 views, 1024×768 detector, pitch 1.25. Default memory settings, without an imposed allocator cap.
 
-To measure one against several GPUs on your machine (default scan of the script, not the table setup):
+Projection times are medians of seven samples after three warmups. CGLS uses three ten-iteration solves after a two-iteration warmup; the median includes initialization and is divided by ten. Timings use synchronized wall time and the slowest rank. Raw samples and configuration: [scaling.json](https://github.com/sypsyp97/diffct/blob/main/docs/assets/scaling.json).
+
+Run the same workload:
 
 ```bash
-python examples/benchmark_projector.py --devices 0 1
+python examples/benchmark_scaling.py --output one.json
+torchrun --standalone --nproc-per-node=4 examples/benchmark_scaling.py --expected-nodes 1 --output four.json
 ```
 
 ## Gallery

{diffct-2.1.0 → diffct-2.2.0}/README.md

@@ -34,8 +34,6 @@
   <a href="#citation">Citation</a>
 </p>
 
-> **Install** with the CUDA extra, `pip install "diffct[cu12]"` or `pip install "diffct[cu13]"`, after PyTorch. A plain `pip install diffct` needs a system CUDA Toolkit.
->
 > **Version 2.0** replaces the circular-orbit API of 1.x with per-view trajectories. Code written for 1.x needs changes; see the [migration notes](https://github.com/sypsyp97/diffct/blob/main/docs/MIGRATION.md).
 > For Apple Silicon, use the [MLX port by Linda-Sophie Schneider](https://github.com/Linda-SophieSchneider/DiffCT-MLX).
 
@@ -44,10 +42,10 @@
 To our knowledge, diffct is the only open-source GPU CT library that combines arbitrary per-view trajectories, autograd through the volume, first-order gradients through the acquisition geometry, and built-in multi-GPU and multi-node execution. We compared it with LEAP, TIGRE, ASTRA/tomosipo, DiffDRR and others in October 2026.
 
 - **Arbitrary trajectories.** Each view has its own source, detector centre and detector axes. Circular, helical, saddle, sinusoidal, random and calibrated scans use the same code. Trajectory tensors with `requires_grad=True` receive geometry gradients for calibration.
-- **Curved detectors.** A `detector_surface(u, v)` callback defines each pixel's position for arc, cylindrical or other parameterized surfaces. Projection and matched backprojection support first-order surface gradients. See [the surface API](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#parameterized-detector-surfaces).
+- **Curved detectors.** `ParameterizedSurface` defines native arc, cylindrical or arbitrary pixels from global coordinates and explicit parameters, sampled in bounded view/pixel batches. Legacy `detector_surface(u, v)` callbacks remain available. Projection and matched backprojection support first-order surface gradients. See [the surface API](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#parameterized-detector-surfaces).
 - **Matched operators.** `project()` and `backproject()` form an exact adjoint pair for the cell-constant Siddon model. Both support PyTorch autograd, with volume/sinogram gradients and Hessian-vector products.
-- **Large volumes.** CPU volumes and sinograms automatically stream spatial tiles and view batches through CUDA. Full arrays and data gradients stay on CPU; tile sizes follow the available device memory and can be overridden. See [chunked execution](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#chunked-execution).
-- **Multi-GPU and multi-node.** Use `devices=[0, 1, 2, 3]` in one process, or one process per GPU with torchrun and NCCL. Views are partitioned. CPU-backed volumes stay in host RAM while GPUs compute tiles; distributed ranks retain replicated-volume semantics. Speedup depends on workload and communication costs.
+- **Large volumes.** Automatic spatial, view and pixel blocks with bounded transfer buffers control working memory. `project_into`/`backproject_into` accept caller-owned tensor or disk stores; tensor autograd keeps complete data on its input device. See [chunked execution](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#chunked-execution).
+- **Multi-GPU and multi-node.** Use `devices=[0, 1, 2, 3]` in one process, or torchrun and NCCL across ranks. Default `partition="views"` shards projections and replicates volumes. `partition="space"` gives each rank an owned volume slab and combines ray contributions. Speedup depends on workload and communication costs.
 - **Analytical helpers.** `diffct.analytical` provides ramp filters (ram-lak, shepp-logan, cosine, hamming, hann), fan, cone and Parker weights, and FBP and FDK backprojection.
 
 Capabilities, limits and isocenter rules: [docs/REFERENCE.md](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#capabilities-and-limits).
@@ -109,7 +107,7 @@ loss.backward()                     # x.grad = A^T (A x - y)
 
 **Geometry gradients.** Set `requires_grad=True` on trajectory tensors before you build `Projector`. See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#geometry-gradients).
 
-**Several GPUs and nodes.** Use `devices=[0, 1, 2, 3]` in one process, or one process per GPU with torchrun for one or more nodes. See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#several-gpus-and-nodes) and [DISTRIBUTED.md](https://github.com/sypsyp97/diffct/blob/main/docs/DISTRIBUTED.md).
+**Several GPUs and nodes.** See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#several-gpus-and-nodes) and [DISTRIBUTED.md](https://github.com/sypsyp97/diffct/blob/main/docs/DISTRIBUTED.md).
 
 ```python
 A = Projector(trajectory, (128, 128, 128), (384, 256), detector_spacing=0.8, devices=[0, 1, 2, 3])
@@ -130,29 +128,36 @@ sinogram.square().mean().backward()          # CPU volume.grad
 Run `python examples/chunked_reconstruction.py` for CPU-backed CGLS and measured
 CUDA allocation. Chunking preserves the Siddon model, with float32 summation
 differences; smaller blocks can increase transfer and launch overhead. Full
-arrays and iterative state still need enough host RAM.
+tensor arrays and iterative state still need enough host RAM. For disk-backed
+state and sharded checkpoints, run
+`python examples/disk_reconstruction.py --output disk-run`; for fixed rank-local
+volume slabs use `--partition space` under torchrun. Run
+`python examples/benchmark_execution.py` to compare schedules and copy payloads.
 
 ## Performance
 
 <p align="center">
   <picture>
-    <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_dark.png">
-    <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_light.png" width="100%" alt="Time per forward, adjoint and CGLS iteration on 1, 4 and 8 GPUs, 64^3 and 128^3">
+    <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_dark.png?v=2.2.0">
+    <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_light.png?v=2.2.0" width="100%" alt="Time per forward, adjoint and CGLS iteration on 1, 4 and 8 GPUs, 512^3 volume">
   </picture>
 </p>
 
-| 128³ volume | 1 GPU | 4 GPUs, 1 node | 8 GPUs, 2 nodes |
+| 512³ volume | 1 GPU | 4 GPUs, 1 node | 8 GPUs, 2 nodes |
 |---|---:|---:|---:|
-| Forward projection | 8.45 ms | 2.66 ms (3.2×) | 1.68 ms (5.0×) |
-| Adjoint (backprojection) | 20.66 ms | 5.89 ms (3.5×) | 4.62 ms (4.5×) |
-| CGLS iteration | 30.08 ms | 8.08 ms (3.7×) | 5.59 ms (5.4×) |
+| Forward projection | 493.05 ms | 132.22 ms (3.7×) | 73.18 ms (6.7×) |
+| Adjoint (backprojection) | 1400.80 ms | 369.29 ms (3.8×) | 265.87 ms (5.3×) |
+| CGLS iteration | 2054.36 ms | 547.06 ms (3.8×) | 373.38 ms (5.5×) |
 
-Circular trajectory, 360 views, detector (2n, 1.5n) cells at pitch 1.25, A100 64 GB GPUs. Small volumes scale less: at 64³, one CGLS iteration takes 4.63 ms on 1 GPU, 1.57 ms on 4 and 1.56 ms on 8. Raw data: [docs/assets/scaling.json](https://github.com/sypsyp97/diffct/blob/main/docs/assets/scaling.json).
+diffct 2.2.0, FP32, A100 64 GB GPUs; circular cone beam, 360 views, 1024×768 detector, pitch 1.25. Default memory settings, without an imposed allocator cap.
 
-To measure one against several GPUs on your machine (default scan of the script, not the table setup):
+Projection times are medians of seven samples after three warmups. CGLS uses three ten-iteration solves after a two-iteration warmup; the median includes initialization and is divided by ten. Timings use synchronized wall time and the slowest rank. Raw samples and configuration: [scaling.json](https://github.com/sypsyp97/diffct/blob/main/docs/assets/scaling.json).
+
+Run the same workload:
 
 ```bash
-python examples/benchmark_projector.py --devices 0 1
+python examples/benchmark_scaling.py --output one.json
+torchrun --standalone --nproc-per-node=4 examples/benchmark_scaling.py --expected-nodes 1 --output four.json
 ```
 
 ## Gallery

{diffct-2.1.0 → diffct-2.2.0}/diffct/__init__.py

@@ -15,6 +15,8 @@ from .projectors import (
 )
 
 from .operators import Projector
+from .storage import TensorStore, NpyStore
+from .surfaces import ParameterizedSurface
 
 from .geometry import (
     circular_trajectory_3d,
@@ -49,10 +51,13 @@ try:
 except ImportError:
     pass
 
-__version__ = '2.1.0'
+__version__ = '2.2.0'
 
 __all__ = [
     'Projector',
+    'TensorStore',
+    'NpyStore',
+    'ParameterizedSurface',
     'ParallelProjectorFunction',
     'ParallelBackprojectorFunction',
     'FanProjectorFunction',

diffct-2.2.0/diffct/chunking.py

@@ -0,0 +1,201 @@
+"""Spatial slices and immutable execution limits for streamed projectors."""
+
+from dataclasses import dataclass
+from itertools import product
+import math
+
+import torch
+
+
+@dataclass(frozen=True)
+class _ChunkPlan:
+    chunk_shape: tuple
+    view_chunk_size: int
+    devices: tuple
+    detector_chunk_shape: tuple = None
+    gpu_budget_bytes: int = 0
+    allocation_limit: int = (1 << 63) - 1
+    schedule: str = "spatial"
+    window_size: int = 1
+    slots: int = 1
+    estimated_gpu_bytes: int = 0
+    fallback_reason: str = ""
+    foreign_output: bool = False
+    reservation_budget_bytes: int = (1 << 63) - 1
+    estimated_reserved_bytes: int = 0
+
+
+def _pixel_tiles(shape, chunk_shape):
+    for starts in product(*(range(0, size, part) for size, part in zip(shape, chunk_shape))):
+        clipped = tuple(min(part, size - start) for size, start, part in zip(shape, starts, chunk_shape))
+        yield tuple(slice(start, start + size) for start, size in zip(starts, clipped)), clipped
+
+
+def _memory_budget(device):
+    """Usable allocator bytes, retaining 25% headroom for allocation overhead.
+
+    Reusable reserved blocks count as available, but never bypass the process
+    fraction ceiling. Older PyTorch without its public fraction getter can
+    only observe driver/allocator memory; explicit limits remain available.
+    """
+    free, total = torch.cuda.mem_get_info(device)
+    allocated = torch.cuda.memory_allocated(device)
+    reusable = max(0, torch.cuda.memory_reserved(device) - allocated)
+    getter = getattr(torch.cuda, "get_per_process_memory_fraction", None)
+    fraction = getter(device) if getter is not None else 1.0
+    allocator_room = max(0, int(total * fraction) - allocated)
+    backend = getattr(torch.cuda, "get_allocator_backend", lambda: "native")()
+    foreign_splits = 0
+    if backend == "native":
+        caller = torch.cuda.current_stream(device).cuda_stream
+        for segment in _native_segments(device):
+            blocks = segment["blocks"]
+            if segment["stream"] != caller and any(block["state"] != "inactive" for block in blocks):
+                foreign_splits += sum(block["size"] for block in blocks if block["state"] == "inactive")
+    available = max(0, min(free + reusable, allocator_room) - foreign_splits)
+    budget = available * 3 // 4
+    allocation_limit = (1 << 63) - 1  # Unbounded limit in scalar int64 agreement.
+    # Native requests >1 MiB and <10 MiB reserve a 20 MiB segment. A smaller
+    # current available room must use the <=1 MiB small pool, whose segments are 2 MiB.
+    # Source: pytorch/c10/cuda/CUDACachingAllocator.cpp (Native allocator notes).
+    if backend == "native" and available < 20 * 1024 ** 2:
+        allocation_limit = 1024 ** 2
+    return budget, allocation_limit
+
+
+def _native_segments(device):
+    index = torch.device(device).index if not isinstance(device, int) else device
+    index = torch.cuda.current_device() if index is None else index
+    return [segment for segment in torch.cuda.memory_snapshot() if segment["device"] == index]
+
+
+def _reservation_budget(device):
+    """Room for fresh segments, conservatively excluding every pinned segment.
+
+    Free splits in a live segment may serve the caller, but cannot guarantee
+    placement of the complete new inventory. Wholly inactive segments can be
+    released by the native allocator, including those owned by other streams.
+    """
+    if torch.cuda.get_allocator_backend() != "native":
+        return (1 << 63) - 1
+    free, total = torch.cuda.mem_get_info(device)
+    segments = _native_segments(device)
+    pinned = sum(segment["total_size"] for segment in segments
+                 if any(block["state"] != "inactive" for block in segment["blocks"]))
+    reclaimable = sum(segment["total_size"] for segment in segments
+                      if all(block["state"] == "inactive" for block in segment["blocks"]))
+    getter = getattr(torch.cuda, "get_per_process_memory_fraction", None)
+    fraction = getter(device) if getter is not None else 1.0
+    return max(0, min(free + reclaimable, int(total * fraction) - pinned))
+
+
+def _working_set_bytes(beam, shape, views, detector_shape, surface):
+    """Conservative native staging/layout, cotangent, geometry and VJP scratch."""
+    rank = len(shape)
+    voxels = math.prod(shape)
+    rays = views * math.prod(detector_shape)
+    geometry = (rank * (views + rays) if surface
+                else rank * (4 if beam == "cone" else 3) * views)
+    # A resident native tile, staging and completed cone layout can coexist.
+    # Two ray buffers cover reusable output/cotangent and transfer scratch; two
+    # geometry buffers cover staged values and their first-order VJP/reduction.
+    tile_buffers = 3 if beam == "cone" else 2
+    return 4 * (tile_buffers * voxels + 2 * rays + 2 * geometry)
+
+
+def _largest_buffer_bytes(beam, shape, views, detector_shape, surface):
+    rank = len(shape)
+    rays = views * math.prod(detector_shape)
+    geometry = (rank * (views + rays) if surface
+                else rank * (4 if beam == "cone" else 3) * views)
+    return 4 * max(math.prod(shape), rays, geometry)
+
+
+def _allocation_bytes(elements):
+    size = 4 * elements
+    if torch.cuda.get_allocator_backend() == "native":
+        return (size + 511) // 512 * 512
+    return size
+
+
+def _reservation_bytes(buffer_sizes):
+    """Fresh, monotone fixed-buffer inventory bound for standard native pools.
+
+    Arbitrary free holes and allocator overrides are outside this model. Large
+    and middle requests share one pool, so mixed inventories use separate fresh
+    segments per request instead of assuming their placement is independent.
+    """
+    sizes = [size for size in buffer_sizes if size > 0]
+    if torch.cuda.get_allocator_backend() != "native":
+        return sum(sizes)
+    mib = 1024 ** 2
+    sizes = [(size + 511) // 512 * 512 for size in sizes]
+    large = [size for size in sizes if size >= 10 * mib]
+    reservation = sum((size + 2 * mib - 1) // (2 * mib) * (2 * mib) for size in large)
+    middle = [size for size in sizes if mib < size < 10 * mib]
+    if large:
+        reservation += len(middle) * 20 * mib
+        middle = []
+    for segment_size, requests in (
+        (2 * mib, [size for size in sizes if size <= mib]),
+        (20 * mib, middle),
+    ):
+        if not requests:
+            continue
+        total = sum(requests)
+        # At a new segment's birth every older segment holds more than
+        # segment_size - max(requests); the new one holds at least min(requests).
+        # This bounds online best-fit placement in any allocation order.
+        count = (1 if total <= segment_size else
+                 (total - min(requests)) // (segment_size - max(requests)) + 1)
+        reservation += count * segment_size
+    return reservation
+
+
+def _pipeline_buffers(beam, shape, views, pixels, points, *, slots=1, retained=1,
+                      geometry_grad=False, peers=False, ray_accumulators=0):
+    """Live typed allocation inventory, with native rounding per buffer."""
+    rank, rays = len(shape), views * math.prod(pixels)
+    native = _allocation_bytes(math.prod(shape))
+    # Two producer epochs permit the preceding native/layout/output storage
+    # to stay in flight while the next retained tile window queues work.
+    tiles = slots * (retained + int(beam == "cone")) + int(peers)
+    sizes = ((views * rank, rays * rank) if points
+             else (views * rank,) * (4 if beam == "cone" else 3))
+    geometry = tuple(_allocation_bytes(size) for size in sizes)
+    return ((native,) * tiles + (_allocation_bytes(rays),) * (slots + ray_accumulators)
+            + geometry * (slots * (1 + int(geometry_grad))))
+
+
+def _pipeline_bytes(beam, shape, views, pixels, points, *, slots=1, retained=1,
+                    geometry_grad=False, peers=False, ray_accumulators=0):
+    """Live typed bytes; segment reservation is admitted separately."""
+    rank, rays = len(shape), views * math.prod(pixels)
+    native = _allocation_bytes(math.prod(shape))
+    tile = native * (slots * (retained + int(beam == "cone")) + int(peers))
+    sizes = ((views * rank, rays * rank) if points
+             else (views * rank,) * (4 if beam == "cone" else 3))
+    geometry = sum(_allocation_bytes(size) for size in sizes)
+    return (tile + slots * (_allocation_bytes(rays) + geometry)
+            + ray_accumulators * _allocation_bytes(rays) + (slots * geometry if geometry_grad else 0))
+
+
+def _host_geometry_bytes(beam, views, pixels, points):
+    rank = 3 if beam == "cone" else 2
+    if not points:
+        return 8 * 2 * rank * views * (4 if beam == "cone" else 3)
+    # Grids/offsets, world expressions and first-order cotangent/graph work.
+    # Arbitrary sampler internals and explicit parameter snapshots are separate.
+    return 8 * rank * (math.prod(pixels) * (1 + 6 * views) + 4 * views)
+
+
+def _tiles(shape, chunk_shape, spacing):
+    """Yield tensor slices, clipped shapes and physical centres in xyz order."""
+    starts = (range(0, size, chunk) for size, chunk in zip(shape, chunk_shape))
+    for start in product(*starts):
+        tile_shape = tuple(min(chunk, size - first)
+                           for size, chunk, first in zip(shape, chunk_shape, start))
+        slices = tuple(slice(first, first + size) for first, size in zip(start, tile_shape))
+        centre = tuple((first + size / 2 - full / 2) * spacing
+                       for first, size, full in zip(start, tile_shape, shape))
+        yield slices, tile_shape, centre[::-1]