diffct 2.0.2__tar.gz → 2.2.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {diffct-2.0.2 → diffct-2.2.0}/CHANGELOG.md +72 -0
- {diffct-2.0.2 → diffct-2.2.0}/CITATION.cff +2 -2
- {diffct-2.0.2 → diffct-2.2.0}/PKG-INFO +45 -19
- {diffct-2.0.2 → diffct-2.2.0}/README.md +44 -18
- {diffct-2.0.2 → diffct-2.2.0}/diffct/__init__.py +6 -1
- diffct-2.2.0/diffct/chunking.py +201 -0
- {diffct-2.0.2 → diffct-2.2.0}/diffct/kernels/cone_beam.py +46 -36
- {diffct-2.0.2 → diffct-2.2.0}/diffct/kernels/fan_beam.py +30 -22
- {diffct-2.0.2 → diffct-2.2.0}/diffct/kernels/geometry_vjp.py +77 -48
- {diffct-2.0.2 → diffct-2.2.0}/diffct/kernels/parallel_beam.py +32 -22
- diffct-2.2.0/diffct/operators.py +2047 -0
- diffct-2.2.0/diffct/pipeline.py +258 -0
- {diffct-2.0.2 → diffct-2.2.0}/diffct/projectors.py +155 -125
- diffct-2.2.0/diffct/storage.py +125 -0
- diffct-2.2.0/diffct/streaming.py +48 -0
- diffct-2.2.0/diffct/surfaces.py +103 -0
- {diffct-2.0.2 → diffct-2.2.0}/diffct/utils.py +5 -0
- diffct-2.2.0/docs/video/README.md +119 -0
- {diffct-2.0.2 → diffct-2.2.0}/examples/README.md +55 -2
- diffct-2.2.0/formal/README.md +151 -0
- {diffct-2.0.2 → diffct-2.2.0}/pyproject.toml +1 -1
- diffct-2.2.0/tests/test_allocator_capacity.py +349 -0
- diffct-2.2.0/tests/test_axis_aligned_chords.py +118 -0
- diffct-2.2.0/tests/test_batched_surfaces.py +705 -0
- diffct-2.2.0/tests/test_block_cgls.py +996 -0
- diffct-2.2.0/tests/test_block_storage.py +539 -0
- diffct-2.2.0/tests/test_cgls_state_admission.py +173 -0
- diffct-2.2.0/tests/test_chunked_projector.py +1298 -0
- diffct-2.2.0/tests/test_completed_transfer_reuse.py +92 -0
- diffct-2.2.0/tests/test_cuda_sink_lifetimes.py +125 -0
- diffct-2.2.0/tests/test_detector_surface_dtypes.py +80 -0
- diffct-2.2.0/tests/test_detector_surface_mapping.py +79 -0
- diffct-2.2.0/tests/test_detector_surface_units.py +62 -0
- diffct-2.2.0/tests/test_detector_surfaces.py +705 -0
- diffct-2.2.0/tests/test_example_psnr.py +75 -0
- diffct-2.2.0/tests/test_execution_buffers.py +613 -0
- diffct-2.2.0/tests/test_execution_schedules.py +582 -0
- diffct-2.2.0/tests/test_output_pipeline_overlap.py +130 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/test_projector_api.py +9 -3
- diffct-2.2.0/tests/test_spatial_partition.py +433 -0
- diffct-2.2.0/tests/test_transfer_pipeline.py +646 -0
- diffct-2.2.0/tests/test_transfer_retirement_lifetimes.py +166 -0
- diffct-2.0.2/diffct/operators.py +0 -555
- diffct-2.0.2/docs/video/README.md +0 -47
- diffct-2.0.2/formal/README.md +0 -82
- {diffct-2.0.2 → diffct-2.2.0}/.gitignore +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/LICENSE +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/NOTICE +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/diffct/analytical.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/diffct/constants.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/diffct/differentiable.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/diffct/geometry.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/diffct/kernels/__init__.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/examples/data/NOTICE +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/pytest.ini +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/__init__.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/_phantoms.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/benchmarks/README.md +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/benchmarks/__init__.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/benchmarks/_common.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/benchmarks/test_bench_cone.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/benchmarks/test_bench_fan.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/benchmarks/test_bench_parallel.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/distributed_projector_check.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/test_adjoint_inner_product.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/test_analytical_regressions.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/test_cone_projector_autograd.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/test_cuda_exit_deallocations.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/test_cuda_smoke.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/test_fbp_fan_accuracy.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/test_fbp_fan_offsets.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/test_fbp_parallel_accuracy.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/test_fdk_cone_accuracy.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/test_fdk_cone_offsets.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/test_geometry_autograd.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/test_geometry_gradients.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/test_gradcheck.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/test_multi_gpu_exit.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/test_projector_degenerate_geometry.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/test_projector_geometry_devices.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/test_ramp_filter_windows.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/test_release_2_0_1.py +0 -0
- {diffct-2.0.2 → diffct-2.2.0}/tests/test_weights.py +0 -0
|
@@ -5,6 +5,78 @@ All notable changes to this project are documented in this file.
|
|
|
5
5
|
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
|
|
6
6
|
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
|
|
7
7
|
|
|
8
|
+
## [Unreleased]
|
|
9
|
+
|
|
10
|
+
## [2.2.0] - 2026-10-11
|
|
11
|
+
|
|
12
|
+
### Added
|
|
13
|
+
|
|
14
|
+
- Operation-specific automatic schedules and slab/block comparisons for
|
|
15
|
+
native projection and matched backprojection, with bounded CUDA pilots
|
|
16
|
+
and per-call execution/copy metadata.
|
|
17
|
+
- Detector-pixel tiling and explicit `ParameterizedSurface` batches using
|
|
18
|
+
global pixel/view coordinates, saved parameter/frame values and local VJP
|
|
19
|
+
reduction. Legacy two-argument surface callbacks remain compatible.
|
|
20
|
+
- Caller-owned numerical `project_into`/`backproject_into` outputs, tensor
|
|
21
|
+
block access and mapped `.npy` stores for data exceeding host RAM.
|
|
22
|
+
- Distributed `partition="space"` with fixed volume slabs, ownership-aware
|
|
23
|
+
autograd/norms and disk CGLS state/checkpoints without full-volume gathering.
|
|
24
|
+
- Independent CUDA tests and bounded formal protocol/ownership verification.
|
|
25
|
+
|
|
26
|
+
### Changed
|
|
27
|
+
|
|
28
|
+
- GPU backprojection accumulators now span all contributing view batches,
|
|
29
|
+
with final transfer after tile completion. Contiguous native volume layouts
|
|
30
|
+
and ray buffers are reused across batches.
|
|
31
|
+
- Transfer execution uses bounded pinned slots, separate streams and events;
|
|
32
|
+
all extra staging/cached GPU buffers are included in capacity estimates.
|
|
33
|
+
Local GPU and NCCL reductions occur before completed-result downloads.
|
|
34
|
+
- Example CGLS norms accumulate FP64 blocks and vector updates reuse buffers.
|
|
35
|
+
The disk example keeps all three volume states plus residual/projection state
|
|
36
|
+
blockwise and retains `q` until the global `alpha` decision.
|
|
37
|
+
|
|
38
|
+
### Fixed
|
|
39
|
+
|
|
40
|
+
- Tensor CGLS ranks agree on CPU fallback before jointly replanning unequal
|
|
41
|
+
spatial shards, preserving collective order.
|
|
42
|
+
- Parameterized detector surfaces use a common arithmetic dtype of at least
|
|
43
|
+
float32, supporting mixed frame dtypes and finite half-precision inputs.
|
|
44
|
+
- CUDA store uploads wait for newly generated blocks on their source stream,
|
|
45
|
+
including transfers within one GPU.
|
|
46
|
+
- Single-slot execution retires prior native computation before allocating
|
|
47
|
+
the next tile, keeping deferred CUDA allocations within its admission bound.
|
|
48
|
+
- Mapped input and output stores reject hard links to the same file before writes.
|
|
49
|
+
- Example PSNR supports CPU reconstruction results with CUDA reference volumes.
|
|
50
|
+
- Deferred CPU output writes retain their submission order across spatial tiles.
|
|
51
|
+
- Coronal slice displays preserve the measured detector row orientation.
|
|
52
|
+
|
|
53
|
+
## [2.1.0] - 2026-10-10
|
|
54
|
+
|
|
55
|
+
### Added
|
|
56
|
+
|
|
57
|
+
- `Projector(..., detector_surface=surface)` samples local parameterized surfaces
|
|
58
|
+
into per-pixel world positions for native parallel, fan and cone projection
|
|
59
|
+
and matched backprojection. Shared or per-view surfaces can contain trainable
|
|
60
|
+
PyTorch parameters. Flat detectors remain the default; analytical helpers
|
|
61
|
+
retain their flat-detector assumptions.
|
|
62
|
+
- CUDA regression tests and CPU formal verification for detector
|
|
63
|
+
surfaces, including bounded production-kernel checks against cell intersections.
|
|
64
|
+
- Automatic CPU-backed spatial tiles and view batches for native `Projector`
|
|
65
|
+
projection, matched backprojection, data autograd/Hessians and first-order
|
|
66
|
+
geometry gradients. CUDA inputs retain the full-volume path when its estimated
|
|
67
|
+
working set fits. Optional `volume_chunk_shape` and `view_chunk_size` override
|
|
68
|
+
automatic limits; full CPU arrays are not staged on CUDA. Several local
|
|
69
|
+
devices compute their view batches concurrently.
|
|
70
|
+
- A CPU-backed CGLS example and chunk-boundary, CUDA memory and
|
|
71
|
+
formal partition/translation/adjoint checks.
|
|
72
|
+
|
|
73
|
+
### Fixed
|
|
74
|
+
|
|
75
|
+
- The exit-time cleanup hook returns early under the CUDA simulator
|
|
76
|
+
(``NUMBA_ENABLE_CUDASIM=1``). Before, each simulator process printed an
|
|
77
|
+
``ImportError`` for ``_PendingDeallocs`` at exit. The simulator loads no
|
|
78
|
+
CUDA modules, so it needs no cleanup.
|
|
79
|
+
|
|
8
80
|
## [2.0.2] - 2026-10-09
|
|
9
81
|
|
|
10
82
|
### Fixed
|
|
@@ -13,8 +13,8 @@ authors:
|
|
|
13
13
|
given-names: Yipeng
|
|
14
14
|
email: yipeng.sun@fau.de
|
|
15
15
|
alias: sypsyp97
|
|
16
|
-
version: 2.0
|
|
17
|
-
date-released: "2026-10-
|
|
16
|
+
version: 2.2.0
|
|
17
|
+
date-released: "2026-10-11"
|
|
18
18
|
license: Apache-2.0
|
|
19
19
|
repository-code: "https://github.com/sypsyp97/diffct"
|
|
20
20
|
url: "https://sypsyp97.github.io/diffct/"
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: diffct
|
|
3
|
-
Version: 2.0
|
|
3
|
+
Version: 2.2.0
|
|
4
4
|
Summary: Differentiable CUDA projectors for CT: arbitrary trajectories, multi-GPU and multi-node execution, geometry gradients
|
|
5
5
|
Project-URL: Homepage, https://github.com/sypsyp97/diffct
|
|
6
6
|
Project-URL: Documentation, https://sypsyp97.github.io/diffct/
|
|
@@ -47,11 +47,11 @@ Description-Content-Type: text/markdown
|
|
|
47
47
|
</p>
|
|
48
48
|
|
|
49
49
|
<p align="center">
|
|
50
|
-
<img src="https://raw.githubusercontent.com/sypsyp97/diffct/
|
|
50
|
+
<img src="https://raw.githubusercontent.com/sypsyp97/diffct/edf7c52227ec2cc14c988918faa7313f8ee0fe03/docs/assets/diffct_intro.gif" width="100%" alt="diffct intro: projection, sinogram, trajectories, multi-GPU split, reconstruction and geometry calibration">
|
|
51
51
|
</p>
|
|
52
52
|
|
|
53
53
|
<p align="center">
|
|
54
|
-
<a href="https://github.com/sypsyp97/diffct/blob/
|
|
54
|
+
<a href="https://github.com/sypsyp97/diffct/blob/edf7c52227ec2cc14c988918faa7313f8ee0fe03/docs/assets/diffct_intro.mp4">Intro video (MP4)</a> ·
|
|
55
55
|
<a href="https://www.preprints.org/manuscript/202605.1446/v1">Technical report</a> ·
|
|
56
56
|
<a href="https://doi.org/10.20944/preprints202605.1446.v1">DOI</a>
|
|
57
57
|
</p>
|
|
@@ -65,8 +65,6 @@ Description-Content-Type: text/markdown
|
|
|
65
65
|
<a href="#citation">Citation</a>
|
|
66
66
|
</p>
|
|
67
67
|
|
|
68
|
-
> **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.
|
|
69
|
-
>
|
|
70
68
|
> **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).
|
|
71
69
|
> For Apple Silicon, use the [MLX port by Linda-Sophie Schneider](https://github.com/Linda-SophieSchneider/DiffCT-MLX).
|
|
72
70
|
|
|
@@ -75,8 +73,10 @@ Description-Content-Type: text/markdown
|
|
|
75
73
|
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.
|
|
76
74
|
|
|
77
75
|
- **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.
|
|
76
|
+
- **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).
|
|
78
77
|
- **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.
|
|
79
|
-
- **
|
|
78
|
+
- **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).
|
|
79
|
+
- **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.
|
|
80
80
|
- **Analytical helpers.** `diffct.analytical` provides ramp filters (ram-lak, shepp-logan, cosine, hamming, hann), fan, cone and Parker weights, and FBP and FDK backprojection.
|
|
81
81
|
|
|
82
82
|
Capabilities, limits and isocenter rules: [docs/REFERENCE.md](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#capabilities-and-limits).
|
|
@@ -138,41 +138,67 @@ loss.backward() # x.grad = A^T (A x - y)
|
|
|
138
138
|
|
|
139
139
|
**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).
|
|
140
140
|
|
|
141
|
-
**Several GPUs and nodes.**
|
|
141
|
+
**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).
|
|
142
142
|
|
|
143
143
|
```python
|
|
144
144
|
A = Projector(trajectory, (128, 128, 128), (384, 256), detector_spacing=0.8, devices=[0, 1, 2, 3])
|
|
145
145
|
```
|
|
146
146
|
|
|
147
|
+
**Volumes larger than GPU memory.** Keep the arrays on CPU. The same calls use
|
|
148
|
+
CUDA tiles automatically; CUDA inputs retain the full-volume path when its
|
|
149
|
+
estimated working set fits. Explicit limits are optional:
|
|
150
|
+
|
|
151
|
+
```python
|
|
152
|
+
A = Projector(trajectory, (128, 128, 128), (384, 256), detector_spacing=0.8)
|
|
153
|
+
volume = torch.ones(A.volume_shape, device="cpu", requires_grad=True)
|
|
154
|
+
sinogram = A.project(volume) # CPU output, CUDA computation
|
|
155
|
+
sinogram.square().mean().backward() # CPU volume.grad
|
|
156
|
+
# Optional: volume_chunk_shape=(32, 64, 64), view_chunk_size=16
|
|
157
|
+
```
|
|
158
|
+
|
|
159
|
+
Run `python examples/chunked_reconstruction.py` for CPU-backed CGLS and measured
|
|
160
|
+
CUDA allocation. Chunking preserves the Siddon model, with float32 summation
|
|
161
|
+
differences; smaller blocks can increase transfer and launch overhead. Full
|
|
162
|
+
tensor arrays and iterative state still need enough host RAM. For disk-backed
|
|
163
|
+
state and sharded checkpoints, run
|
|
164
|
+
`python examples/disk_reconstruction.py --output disk-run`; for fixed rank-local
|
|
165
|
+
volume slabs use `--partition space` under torchrun. Run
|
|
166
|
+
`python examples/benchmark_execution.py` to compare schedules and copy payloads.
|
|
167
|
+
|
|
147
168
|
## Performance
|
|
148
169
|
|
|
149
170
|
<p align="center">
|
|
150
171
|
<picture>
|
|
151
|
-
<source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_dark.png">
|
|
152
|
-
<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,
|
|
172
|
+
<source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_dark.png?v=2.2.0">
|
|
173
|
+
<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">
|
|
153
174
|
</picture>
|
|
154
175
|
</p>
|
|
155
176
|
|
|
156
|
-
|
|
|
177
|
+
| 512³ volume | 1 GPU | 4 GPUs, 1 node | 8 GPUs, 2 nodes |
|
|
157
178
|
|---|---:|---:|---:|
|
|
158
|
-
| Forward projection |
|
|
159
|
-
| Adjoint (backprojection) |
|
|
160
|
-
| CGLS iteration |
|
|
179
|
+
| Forward projection | 493.05 ms | 132.22 ms (3.7×) | 73.18 ms (6.7×) |
|
|
180
|
+
| Adjoint (backprojection) | 1400.80 ms | 369.29 ms (3.8×) | 265.87 ms (5.3×) |
|
|
181
|
+
| CGLS iteration | 2054.36 ms | 547.06 ms (3.8×) | 373.38 ms (5.5×) |
|
|
182
|
+
|
|
183
|
+
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.
|
|
161
184
|
|
|
162
|
-
|
|
185
|
+
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).
|
|
163
186
|
|
|
164
|
-
|
|
187
|
+
Run the same workload:
|
|
165
188
|
|
|
166
189
|
```bash
|
|
167
|
-
python examples/
|
|
190
|
+
python examples/benchmark_scaling.py --output one.json
|
|
191
|
+
torchrun --standalone --nproc-per-node=4 examples/benchmark_scaling.py --expected-nodes 1 --output four.json
|
|
168
192
|
```
|
|
169
193
|
|
|
170
194
|
## Gallery
|
|
171
195
|
|
|
196
|
+
TV denotes TV-regularized iterative reconstruction; Adam is the optimizer. Figure labels distinguish the method from its optimizer.
|
|
197
|
+
|
|
172
198
|
**Measured walnut.** 240 measured views, circular cone beam, 256³: FDK (Hann window), SIRT (200 iterations), CGLS (20), TV (300, weight 0.3); axial and coronal centre slices.
|
|
173
199
|
|
|
174
200
|
<p align="center">
|
|
175
|
-
<img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/walnut_measured.png" width="100%" alt="Measured walnut: FDK, SIRT, CGLS and TV reconstructions">
|
|
201
|
+
<img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/walnut_measured.png?v=e0092c4" width="100%" alt="Measured walnut: FDK, SIRT, CGLS and TV reconstructions">
|
|
176
202
|
</p>
|
|
177
203
|
|
|
178
204
|
```bash
|
|
@@ -182,7 +208,7 @@ python examples/walnut_reconstruction.py --figure out.png
|
|
|
182
208
|
**Simulated helical scan of the walnut.** The FDK walnut volume as ground truth, 720 views, 1% noise, 256³. PSNR: FDK 26.67 dB, CGLS (30 iterations) 34.13 dB, SIRT (200) 34.52 dB, TV (200, weight 1.0) 37.77 dB.
|
|
183
209
|
|
|
184
210
|
<p align="center">
|
|
185
|
-
<img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/walnut_helical.png" width="100%" alt="Simulated helical scan of the walnut: FDK, CGLS, SIRT and TV reconstructions">
|
|
211
|
+
<img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/walnut_helical.png?v=e0092c4" width="100%" alt="Simulated helical scan of the walnut: FDK, CGLS, SIRT and TV reconstructions">
|
|
186
212
|
</p>
|
|
187
213
|
|
|
188
214
|
```bash
|
|
@@ -194,7 +220,7 @@ Walnut data: Meaney 2022, Zenodo 6986012, CC BY 4.0; see [examples/data/NOTICE](
|
|
|
194
220
|
|
|
195
221
|
## Examples
|
|
196
222
|
|
|
197
|
-
`quickstart.py`, `analytical_reconstruction.py`, `iterative_reconstruction.py`, `walnut_reconstruction.py`, `geometry_calibration.py`, `benchmark_projector.py`, `plot_trajectory.py`, and the Slurm template `slurm/multi_node.sbatch`. Launch modes and distributed-loss rules are in [examples/README.md](https://github.com/sypsyp97/diffct/blob/main/examples/README.md).
|
|
223
|
+
`quickstart.py`, `curved_detector.py`, `chunked_reconstruction.py`, `analytical_reconstruction.py`, `iterative_reconstruction.py`, `walnut_reconstruction.py`, `geometry_calibration.py`, `benchmark_projector.py`, `plot_trajectory.py`, and the Slurm template `slurm/multi_node.sbatch`. Launch modes and distributed-loss rules are in [examples/README.md](https://github.com/sypsyp97/diffct/blob/main/examples/README.md).
|
|
198
224
|
|
|
199
225
|
## Documentation
|
|
200
226
|
|
|
@@ -16,11 +16,11 @@
|
|
|
16
16
|
</p>
|
|
17
17
|
|
|
18
18
|
<p align="center">
|
|
19
|
-
<img src="https://raw.githubusercontent.com/sypsyp97/diffct/
|
|
19
|
+
<img src="https://raw.githubusercontent.com/sypsyp97/diffct/edf7c52227ec2cc14c988918faa7313f8ee0fe03/docs/assets/diffct_intro.gif" width="100%" alt="diffct intro: projection, sinogram, trajectories, multi-GPU split, reconstruction and geometry calibration">
|
|
20
20
|
</p>
|
|
21
21
|
|
|
22
22
|
<p align="center">
|
|
23
|
-
<a href="https://github.com/sypsyp97/diffct/blob/
|
|
23
|
+
<a href="https://github.com/sypsyp97/diffct/blob/edf7c52227ec2cc14c988918faa7313f8ee0fe03/docs/assets/diffct_intro.mp4">Intro video (MP4)</a> ·
|
|
24
24
|
<a href="https://www.preprints.org/manuscript/202605.1446/v1">Technical report</a> ·
|
|
25
25
|
<a href="https://doi.org/10.20944/preprints202605.1446.v1">DOI</a>
|
|
26
26
|
</p>
|
|
@@ -34,8 +34,6 @@
|
|
|
34
34
|
<a href="#citation">Citation</a>
|
|
35
35
|
</p>
|
|
36
36
|
|
|
37
|
-
> **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.
|
|
38
|
-
>
|
|
39
37
|
> **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).
|
|
40
38
|
> For Apple Silicon, use the [MLX port by Linda-Sophie Schneider](https://github.com/Linda-SophieSchneider/DiffCT-MLX).
|
|
41
39
|
|
|
@@ -44,8 +42,10 @@
|
|
|
44
42
|
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.
|
|
45
43
|
|
|
46
44
|
- **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.
|
|
45
|
+
- **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).
|
|
47
46
|
- **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.
|
|
48
|
-
- **
|
|
47
|
+
- **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).
|
|
48
|
+
- **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.
|
|
49
49
|
- **Analytical helpers.** `diffct.analytical` provides ramp filters (ram-lak, shepp-logan, cosine, hamming, hann), fan, cone and Parker weights, and FBP and FDK backprojection.
|
|
50
50
|
|
|
51
51
|
Capabilities, limits and isocenter rules: [docs/REFERENCE.md](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#capabilities-and-limits).
|
|
@@ -107,41 +107,67 @@ loss.backward() # x.grad = A^T (A x - y)
|
|
|
107
107
|
|
|
108
108
|
**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).
|
|
109
109
|
|
|
110
|
-
**Several GPUs and nodes.**
|
|
110
|
+
**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).
|
|
111
111
|
|
|
112
112
|
```python
|
|
113
113
|
A = Projector(trajectory, (128, 128, 128), (384, 256), detector_spacing=0.8, devices=[0, 1, 2, 3])
|
|
114
114
|
```
|
|
115
115
|
|
|
116
|
+
**Volumes larger than GPU memory.** Keep the arrays on CPU. The same calls use
|
|
117
|
+
CUDA tiles automatically; CUDA inputs retain the full-volume path when its
|
|
118
|
+
estimated working set fits. Explicit limits are optional:
|
|
119
|
+
|
|
120
|
+
```python
|
|
121
|
+
A = Projector(trajectory, (128, 128, 128), (384, 256), detector_spacing=0.8)
|
|
122
|
+
volume = torch.ones(A.volume_shape, device="cpu", requires_grad=True)
|
|
123
|
+
sinogram = A.project(volume) # CPU output, CUDA computation
|
|
124
|
+
sinogram.square().mean().backward() # CPU volume.grad
|
|
125
|
+
# Optional: volume_chunk_shape=(32, 64, 64), view_chunk_size=16
|
|
126
|
+
```
|
|
127
|
+
|
|
128
|
+
Run `python examples/chunked_reconstruction.py` for CPU-backed CGLS and measured
|
|
129
|
+
CUDA allocation. Chunking preserves the Siddon model, with float32 summation
|
|
130
|
+
differences; smaller blocks can increase transfer and launch overhead. Full
|
|
131
|
+
tensor arrays and iterative state still need enough host RAM. For disk-backed
|
|
132
|
+
state and sharded checkpoints, run
|
|
133
|
+
`python examples/disk_reconstruction.py --output disk-run`; for fixed rank-local
|
|
134
|
+
volume slabs use `--partition space` under torchrun. Run
|
|
135
|
+
`python examples/benchmark_execution.py` to compare schedules and copy payloads.
|
|
136
|
+
|
|
116
137
|
## Performance
|
|
117
138
|
|
|
118
139
|
<p align="center">
|
|
119
140
|
<picture>
|
|
120
|
-
<source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_dark.png">
|
|
121
|
-
<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,
|
|
141
|
+
<source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_dark.png?v=2.2.0">
|
|
142
|
+
<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">
|
|
122
143
|
</picture>
|
|
123
144
|
</p>
|
|
124
145
|
|
|
125
|
-
|
|
|
146
|
+
| 512³ volume | 1 GPU | 4 GPUs, 1 node | 8 GPUs, 2 nodes |
|
|
126
147
|
|---|---:|---:|---:|
|
|
127
|
-
| Forward projection |
|
|
128
|
-
| Adjoint (backprojection) |
|
|
129
|
-
| CGLS iteration |
|
|
148
|
+
| Forward projection | 493.05 ms | 132.22 ms (3.7×) | 73.18 ms (6.7×) |
|
|
149
|
+
| Adjoint (backprojection) | 1400.80 ms | 369.29 ms (3.8×) | 265.87 ms (5.3×) |
|
|
150
|
+
| CGLS iteration | 2054.36 ms | 547.06 ms (3.8×) | 373.38 ms (5.5×) |
|
|
151
|
+
|
|
152
|
+
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.
|
|
130
153
|
|
|
131
|
-
|
|
154
|
+
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).
|
|
132
155
|
|
|
133
|
-
|
|
156
|
+
Run the same workload:
|
|
134
157
|
|
|
135
158
|
```bash
|
|
136
|
-
python examples/
|
|
159
|
+
python examples/benchmark_scaling.py --output one.json
|
|
160
|
+
torchrun --standalone --nproc-per-node=4 examples/benchmark_scaling.py --expected-nodes 1 --output four.json
|
|
137
161
|
```
|
|
138
162
|
|
|
139
163
|
## Gallery
|
|
140
164
|
|
|
165
|
+
TV denotes TV-regularized iterative reconstruction; Adam is the optimizer. Figure labels distinguish the method from its optimizer.
|
|
166
|
+
|
|
141
167
|
**Measured walnut.** 240 measured views, circular cone beam, 256³: FDK (Hann window), SIRT (200 iterations), CGLS (20), TV (300, weight 0.3); axial and coronal centre slices.
|
|
142
168
|
|
|
143
169
|
<p align="center">
|
|
144
|
-
<img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/walnut_measured.png" width="100%" alt="Measured walnut: FDK, SIRT, CGLS and TV reconstructions">
|
|
170
|
+
<img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/walnut_measured.png?v=e0092c4" width="100%" alt="Measured walnut: FDK, SIRT, CGLS and TV reconstructions">
|
|
145
171
|
</p>
|
|
146
172
|
|
|
147
173
|
```bash
|
|
@@ -151,7 +177,7 @@ python examples/walnut_reconstruction.py --figure out.png
|
|
|
151
177
|
**Simulated helical scan of the walnut.** The FDK walnut volume as ground truth, 720 views, 1% noise, 256³. PSNR: FDK 26.67 dB, CGLS (30 iterations) 34.13 dB, SIRT (200) 34.52 dB, TV (200, weight 1.0) 37.77 dB.
|
|
152
178
|
|
|
153
179
|
<p align="center">
|
|
154
|
-
<img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/walnut_helical.png" width="100%" alt="Simulated helical scan of the walnut: FDK, CGLS, SIRT and TV reconstructions">
|
|
180
|
+
<img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/walnut_helical.png?v=e0092c4" width="100%" alt="Simulated helical scan of the walnut: FDK, CGLS, SIRT and TV reconstructions">
|
|
155
181
|
</p>
|
|
156
182
|
|
|
157
183
|
```bash
|
|
@@ -163,7 +189,7 @@ Walnut data: Meaney 2022, Zenodo 6986012, CC BY 4.0; see [examples/data/NOTICE](
|
|
|
163
189
|
|
|
164
190
|
## Examples
|
|
165
191
|
|
|
166
|
-
`quickstart.py`, `analytical_reconstruction.py`, `iterative_reconstruction.py`, `walnut_reconstruction.py`, `geometry_calibration.py`, `benchmark_projector.py`, `plot_trajectory.py`, and the Slurm template `slurm/multi_node.sbatch`. Launch modes and distributed-loss rules are in [examples/README.md](https://github.com/sypsyp97/diffct/blob/main/examples/README.md).
|
|
192
|
+
`quickstart.py`, `curved_detector.py`, `chunked_reconstruction.py`, `analytical_reconstruction.py`, `iterative_reconstruction.py`, `walnut_reconstruction.py`, `geometry_calibration.py`, `benchmark_projector.py`, `plot_trajectory.py`, and the Slurm template `slurm/multi_node.sbatch`. Launch modes and distributed-loss rules are in [examples/README.md](https://github.com/sypsyp97/diffct/blob/main/examples/README.md).
|
|
167
193
|
|
|
168
194
|
## Documentation
|
|
169
195
|
|
|
@@ -15,6 +15,8 @@ from .projectors import (
|
|
|
15
15
|
)
|
|
16
16
|
|
|
17
17
|
from .operators import Projector
|
|
18
|
+
from .storage import TensorStore, NpyStore
|
|
19
|
+
from .surfaces import ParameterizedSurface
|
|
18
20
|
|
|
19
21
|
from .geometry import (
|
|
20
22
|
circular_trajectory_3d,
|
|
@@ -49,10 +51,13 @@ try:
|
|
|
49
51
|
except ImportError:
|
|
50
52
|
pass
|
|
51
53
|
|
|
52
|
-
__version__ = '2.0
|
|
54
|
+
__version__ = '2.2.0'
|
|
53
55
|
|
|
54
56
|
__all__ = [
|
|
55
57
|
'Projector',
|
|
58
|
+
'TensorStore',
|
|
59
|
+
'NpyStore',
|
|
60
|
+
'ParameterizedSurface',
|
|
56
61
|
'ParallelProjectorFunction',
|
|
57
62
|
'ParallelBackprojectorFunction',
|
|
58
63
|
'FanProjectorFunction',
|
|
@@ -0,0 +1,201 @@
|
|
|
1
|
+
"""Spatial slices and immutable execution limits for streamed projectors."""
|
|
2
|
+
|
|
3
|
+
from dataclasses import dataclass
|
|
4
|
+
from itertools import product
|
|
5
|
+
import math
|
|
6
|
+
|
|
7
|
+
import torch
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
@dataclass(frozen=True)
|
|
11
|
+
class _ChunkPlan:
|
|
12
|
+
chunk_shape: tuple
|
|
13
|
+
view_chunk_size: int
|
|
14
|
+
devices: tuple
|
|
15
|
+
detector_chunk_shape: tuple = None
|
|
16
|
+
gpu_budget_bytes: int = 0
|
|
17
|
+
allocation_limit: int = (1 << 63) - 1
|
|
18
|
+
schedule: str = "spatial"
|
|
19
|
+
window_size: int = 1
|
|
20
|
+
slots: int = 1
|
|
21
|
+
estimated_gpu_bytes: int = 0
|
|
22
|
+
fallback_reason: str = ""
|
|
23
|
+
foreign_output: bool = False
|
|
24
|
+
reservation_budget_bytes: int = (1 << 63) - 1
|
|
25
|
+
estimated_reserved_bytes: int = 0
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
def _pixel_tiles(shape, chunk_shape):
|
|
29
|
+
for starts in product(*(range(0, size, part) for size, part in zip(shape, chunk_shape))):
|
|
30
|
+
clipped = tuple(min(part, size - start) for size, start, part in zip(shape, starts, chunk_shape))
|
|
31
|
+
yield tuple(slice(start, start + size) for start, size in zip(starts, clipped)), clipped
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def _memory_budget(device):
|
|
35
|
+
"""Usable allocator bytes, retaining 25% headroom for allocation overhead.
|
|
36
|
+
|
|
37
|
+
Reusable reserved blocks count as available, but never bypass the process
|
|
38
|
+
fraction ceiling. Older PyTorch without its public fraction getter can
|
|
39
|
+
only observe driver/allocator memory; explicit limits remain available.
|
|
40
|
+
"""
|
|
41
|
+
free, total = torch.cuda.mem_get_info(device)
|
|
42
|
+
allocated = torch.cuda.memory_allocated(device)
|
|
43
|
+
reusable = max(0, torch.cuda.memory_reserved(device) - allocated)
|
|
44
|
+
getter = getattr(torch.cuda, "get_per_process_memory_fraction", None)
|
|
45
|
+
fraction = getter(device) if getter is not None else 1.0
|
|
46
|
+
allocator_room = max(0, int(total * fraction) - allocated)
|
|
47
|
+
backend = getattr(torch.cuda, "get_allocator_backend", lambda: "native")()
|
|
48
|
+
foreign_splits = 0
|
|
49
|
+
if backend == "native":
|
|
50
|
+
caller = torch.cuda.current_stream(device).cuda_stream
|
|
51
|
+
for segment in _native_segments(device):
|
|
52
|
+
blocks = segment["blocks"]
|
|
53
|
+
if segment["stream"] != caller and any(block["state"] != "inactive" for block in blocks):
|
|
54
|
+
foreign_splits += sum(block["size"] for block in blocks if block["state"] == "inactive")
|
|
55
|
+
available = max(0, min(free + reusable, allocator_room) - foreign_splits)
|
|
56
|
+
budget = available * 3 // 4
|
|
57
|
+
allocation_limit = (1 << 63) - 1 # Unbounded limit in scalar int64 agreement.
|
|
58
|
+
# Native requests >1 MiB and <10 MiB reserve a 20 MiB segment. A smaller
|
|
59
|
+
# current available room must use the <=1 MiB small pool, whose segments are 2 MiB.
|
|
60
|
+
# Source: pytorch/c10/cuda/CUDACachingAllocator.cpp (Native allocator notes).
|
|
61
|
+
if backend == "native" and available < 20 * 1024 ** 2:
|
|
62
|
+
allocation_limit = 1024 ** 2
|
|
63
|
+
return budget, allocation_limit
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
def _native_segments(device):
|
|
67
|
+
index = torch.device(device).index if not isinstance(device, int) else device
|
|
68
|
+
index = torch.cuda.current_device() if index is None else index
|
|
69
|
+
return [segment for segment in torch.cuda.memory_snapshot() if segment["device"] == index]
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
def _reservation_budget(device):
|
|
73
|
+
"""Room for fresh segments, conservatively excluding every pinned segment.
|
|
74
|
+
|
|
75
|
+
Free splits in a live segment may serve the caller, but cannot guarantee
|
|
76
|
+
placement of the complete new inventory. Wholly inactive segments can be
|
|
77
|
+
released by the native allocator, including those owned by other streams.
|
|
78
|
+
"""
|
|
79
|
+
if torch.cuda.get_allocator_backend() != "native":
|
|
80
|
+
return (1 << 63) - 1
|
|
81
|
+
free, total = torch.cuda.mem_get_info(device)
|
|
82
|
+
segments = _native_segments(device)
|
|
83
|
+
pinned = sum(segment["total_size"] for segment in segments
|
|
84
|
+
if any(block["state"] != "inactive" for block in segment["blocks"]))
|
|
85
|
+
reclaimable = sum(segment["total_size"] for segment in segments
|
|
86
|
+
if all(block["state"] == "inactive" for block in segment["blocks"]))
|
|
87
|
+
getter = getattr(torch.cuda, "get_per_process_memory_fraction", None)
|
|
88
|
+
fraction = getter(device) if getter is not None else 1.0
|
|
89
|
+
return max(0, min(free + reclaimable, int(total * fraction) - pinned))
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
def _working_set_bytes(beam, shape, views, detector_shape, surface):
|
|
93
|
+
"""Conservative native staging/layout, cotangent, geometry and VJP scratch."""
|
|
94
|
+
rank = len(shape)
|
|
95
|
+
voxels = math.prod(shape)
|
|
96
|
+
rays = views * math.prod(detector_shape)
|
|
97
|
+
geometry = (rank * (views + rays) if surface
|
|
98
|
+
else rank * (4 if beam == "cone" else 3) * views)
|
|
99
|
+
# A resident native tile, staging and completed cone layout can coexist.
|
|
100
|
+
# Two ray buffers cover reusable output/cotangent and transfer scratch; two
|
|
101
|
+
# geometry buffers cover staged values and their first-order VJP/reduction.
|
|
102
|
+
tile_buffers = 3 if beam == "cone" else 2
|
|
103
|
+
return 4 * (tile_buffers * voxels + 2 * rays + 2 * geometry)
|
|
104
|
+
|
|
105
|
+
|
|
106
|
+
def _largest_buffer_bytes(beam, shape, views, detector_shape, surface):
|
|
107
|
+
rank = len(shape)
|
|
108
|
+
rays = views * math.prod(detector_shape)
|
|
109
|
+
geometry = (rank * (views + rays) if surface
|
|
110
|
+
else rank * (4 if beam == "cone" else 3) * views)
|
|
111
|
+
return 4 * max(math.prod(shape), rays, geometry)
|
|
112
|
+
|
|
113
|
+
|
|
114
|
+
def _allocation_bytes(elements):
|
|
115
|
+
size = 4 * elements
|
|
116
|
+
if torch.cuda.get_allocator_backend() == "native":
|
|
117
|
+
return (size + 511) // 512 * 512
|
|
118
|
+
return size
|
|
119
|
+
|
|
120
|
+
|
|
121
|
+
def _reservation_bytes(buffer_sizes):
|
|
122
|
+
"""Fresh, monotone fixed-buffer inventory bound for standard native pools.
|
|
123
|
+
|
|
124
|
+
Arbitrary free holes and allocator overrides are outside this model. Large
|
|
125
|
+
and middle requests share one pool, so mixed inventories use separate fresh
|
|
126
|
+
segments per request instead of assuming their placement is independent.
|
|
127
|
+
"""
|
|
128
|
+
sizes = [size for size in buffer_sizes if size > 0]
|
|
129
|
+
if torch.cuda.get_allocator_backend() != "native":
|
|
130
|
+
return sum(sizes)
|
|
131
|
+
mib = 1024 ** 2
|
|
132
|
+
sizes = [(size + 511) // 512 * 512 for size in sizes]
|
|
133
|
+
large = [size for size in sizes if size >= 10 * mib]
|
|
134
|
+
reservation = sum((size + 2 * mib - 1) // (2 * mib) * (2 * mib) for size in large)
|
|
135
|
+
middle = [size for size in sizes if mib < size < 10 * mib]
|
|
136
|
+
if large:
|
|
137
|
+
reservation += len(middle) * 20 * mib
|
|
138
|
+
middle = []
|
|
139
|
+
for segment_size, requests in (
|
|
140
|
+
(2 * mib, [size for size in sizes if size <= mib]),
|
|
141
|
+
(20 * mib, middle),
|
|
142
|
+
):
|
|
143
|
+
if not requests:
|
|
144
|
+
continue
|
|
145
|
+
total = sum(requests)
|
|
146
|
+
# At a new segment's birth every older segment holds more than
|
|
147
|
+
# segment_size - max(requests); the new one holds at least min(requests).
|
|
148
|
+
# This bounds online best-fit placement in any allocation order.
|
|
149
|
+
count = (1 if total <= segment_size else
|
|
150
|
+
(total - min(requests)) // (segment_size - max(requests)) + 1)
|
|
151
|
+
reservation += count * segment_size
|
|
152
|
+
return reservation
|
|
153
|
+
|
|
154
|
+
|
|
155
|
+
def _pipeline_buffers(beam, shape, views, pixels, points, *, slots=1, retained=1,
|
|
156
|
+
geometry_grad=False, peers=False, ray_accumulators=0):
|
|
157
|
+
"""Live typed allocation inventory, with native rounding per buffer."""
|
|
158
|
+
rank, rays = len(shape), views * math.prod(pixels)
|
|
159
|
+
native = _allocation_bytes(math.prod(shape))
|
|
160
|
+
# Two producer epochs permit the preceding native/layout/output storage
|
|
161
|
+
# to stay in flight while the next retained tile window queues work.
|
|
162
|
+
tiles = slots * (retained + int(beam == "cone")) + int(peers)
|
|
163
|
+
sizes = ((views * rank, rays * rank) if points
|
|
164
|
+
else (views * rank,) * (4 if beam == "cone" else 3))
|
|
165
|
+
geometry = tuple(_allocation_bytes(size) for size in sizes)
|
|
166
|
+
return ((native,) * tiles + (_allocation_bytes(rays),) * (slots + ray_accumulators)
|
|
167
|
+
+ geometry * (slots * (1 + int(geometry_grad))))
|
|
168
|
+
|
|
169
|
+
|
|
170
|
+
def _pipeline_bytes(beam, shape, views, pixels, points, *, slots=1, retained=1,
|
|
171
|
+
geometry_grad=False, peers=False, ray_accumulators=0):
|
|
172
|
+
"""Live typed bytes; segment reservation is admitted separately."""
|
|
173
|
+
rank, rays = len(shape), views * math.prod(pixels)
|
|
174
|
+
native = _allocation_bytes(math.prod(shape))
|
|
175
|
+
tile = native * (slots * (retained + int(beam == "cone")) + int(peers))
|
|
176
|
+
sizes = ((views * rank, rays * rank) if points
|
|
177
|
+
else (views * rank,) * (4 if beam == "cone" else 3))
|
|
178
|
+
geometry = sum(_allocation_bytes(size) for size in sizes)
|
|
179
|
+
return (tile + slots * (_allocation_bytes(rays) + geometry)
|
|
180
|
+
+ ray_accumulators * _allocation_bytes(rays) + (slots * geometry if geometry_grad else 0))
|
|
181
|
+
|
|
182
|
+
|
|
183
|
+
def _host_geometry_bytes(beam, views, pixels, points):
|
|
184
|
+
rank = 3 if beam == "cone" else 2
|
|
185
|
+
if not points:
|
|
186
|
+
return 8 * 2 * rank * views * (4 if beam == "cone" else 3)
|
|
187
|
+
# Grids/offsets, world expressions and first-order cotangent/graph work.
|
|
188
|
+
# Arbitrary sampler internals and explicit parameter snapshots are separate.
|
|
189
|
+
return 8 * rank * (math.prod(pixels) * (1 + 6 * views) + 4 * views)
|
|
190
|
+
|
|
191
|
+
|
|
192
|
+
def _tiles(shape, chunk_shape, spacing):
|
|
193
|
+
"""Yield tensor slices, clipped shapes and physical centres in xyz order."""
|
|
194
|
+
starts = (range(0, size, chunk) for size, chunk in zip(shape, chunk_shape))
|
|
195
|
+
for start in product(*starts):
|
|
196
|
+
tile_shape = tuple(min(chunk, size - first)
|
|
197
|
+
for size, chunk, first in zip(shape, chunk_shape, start))
|
|
198
|
+
slices = tuple(slice(first, first + size) for first, size in zip(start, tile_shape))
|
|
199
|
+
centre = tuple((first + size / 2 - full / 2) * spacing
|
|
200
|
+
for first, size, full in zip(start, tile_shape, shape))
|
|
201
|
+
yield slices, tile_shape, centre[::-1]
|