diffct 2.4.1__tar.gz → 2.4.2__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.4.1 → diffct-2.4.2}/CHANGELOG.md +8 -0
- {diffct-2.4.1 → diffct-2.4.2}/CITATION.cff +2 -2
- {diffct-2.4.1 → diffct-2.4.2}/PKG-INFO +1 -1
- {diffct-2.4.1 → diffct-2.4.2}/diffct/__init__.py +1 -1
- {diffct-2.4.1 → diffct-2.4.2}/diffct/kernels/surface_geometry.py +13 -7
- {diffct-2.4.1 → diffct-2.4.2}/diffct/operators.py +102 -8
- {diffct-2.4.1 → diffct-2.4.2}/diffct/surfaces.py +2 -0
- {diffct-2.4.1 → diffct-2.4.2}/docs/video/README.md +2 -2
- {diffct-2.4.1 → diffct-2.4.2}/pyproject.toml +1 -1
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_surface_geometry_cudasim.py +107 -16
- diffct-2.4.2/tests/test_surface_host_sampling_budget.py +257 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_surface_pipeline_contracts.py +20 -0
- {diffct-2.4.1 → diffct-2.4.2}/.gitignore +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/LICENSE +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/NOTICE +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/README.md +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/diffct/_solver_state.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/diffct/_solver_tv.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/diffct/analytical.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/diffct/analytical_execution.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/diffct/backends.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/diffct/chunking.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/diffct/constants.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/diffct/differentiable.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/diffct/distribution.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/diffct/geometry.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/diffct/kernels/__init__.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/diffct/kernels/analytical_transpose.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/diffct/kernels/cone_beam.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/diffct/kernels/fan_beam.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/diffct/kernels/geometry_vjp.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/diffct/kernels/parallel_beam.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/diffct/pipeline.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/diffct/projectors.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/diffct/reconstruction.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/diffct/solvers.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/diffct/storage.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/diffct/streaming.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/diffct/surface_gradients.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/diffct/utils.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/examples/README.md +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/examples/_block_cgls.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/examples/_common.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/examples/data/NOTICE +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/examples/data/preprocess_walnut.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/examples/walnut_reconstruction.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/formal/README.md +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/pytest.ini +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/__init__.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/_phantoms.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/benchmarks/README.md +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/benchmarks/__init__.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/benchmarks/_common.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/benchmarks/test_bench_cone.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/benchmarks/test_bench_fan.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/benchmarks/test_bench_parallel.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/distributed_projector_check.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_adjoint_inner_product.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_allocator_capacity.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_analytical_autograd.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_analytical_distribution.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_analytical_regressions.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_analytical_weight_distribution.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_angular_low_precision.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_axis_aligned_chords.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_batch_execution_edges.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_batched_surfaces.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_block_cgls.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_block_storage.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_call_trajectory.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_cgls_state_admission.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_checkpoint_metadata_validation.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_chunked_projector.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_completed_transfer_reuse.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_cone_backprojection_transpose.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_cone_projector_autograd.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_cpu_install.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_cuda_exit_deallocations.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_cuda_sink_lifetimes.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_cuda_smoke.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_detector_surface_dtypes.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_detector_surface_mapping.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_detector_surface_units.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_detector_surfaces.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_distributed_training_transforms.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_example_psnr.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_execution_buffers.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_execution_features.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_execution_schedules.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_fbp_fan_accuracy.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_fbp_fan_offsets.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_fbp_parallel_accuracy.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_fdk_cone_accuracy.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_fdk_cone_offsets.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_geometry_autograd.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_geometry_conversions.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_geometry_gradients.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_gradcheck.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_lazy_native_observers.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_live_trajectory_validation.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_multi_gpu_exit.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_native_endpoint_precision.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_native_surface_transforms.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_native_volume_layout.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_native_volume_layout_distributed.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_output_pipeline_overlap.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_parallel3d.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_parallel3d_surface_oracle.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_partition_lifecycle.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_plan_memory_bounds.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_portable_distributed_geometry.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_portable_endpoint_precision.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_portable_execution_bounds.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_portable_numeric_regressions.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_portable_sampler_snapshot.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_portable_surface_admission.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_projector_api.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_projector_degenerate_geometry.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_projector_geometry_devices.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_ramp_filter_windows.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_reconstruction_api.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_reference_backend.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_release_2_0_1.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_review_230.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_review_regressions.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_solver_admission_regressions.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_solver_api.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_solver_checkpoint_edges.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_solver_collective_failures.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_spatial_partition.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_surface_collective_failures.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_surface_composition_kernel.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_surface_compute_shards.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_surface_gpu_execution.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_surface_gradient_replay.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_surface_native_regressions.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_surface_parameter_updates.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_surface_trajectory_transform.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_training_transforms.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_trajectory_admission_cost.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_trajectory_surface_integration.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_transfer_pipeline.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_transfer_retirement_lifetimes.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_tv_zero_dependency.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_two_far_endpoints.py +0 -0
- {diffct-2.4.1 → diffct-2.4.2}/tests/test_weights.py +0 -0
|
@@ -7,6 +7,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
|
|
7
7
|
|
|
8
8
|
## [Unreleased]
|
|
9
9
|
|
|
10
|
+
## [2.4.2] - 2026-10-21
|
|
11
|
+
|
|
12
|
+
### Fixed
|
|
13
|
+
|
|
14
|
+
- Count host world points in the streamed geometry budget only for surfaces that expand them on the host. Shared, non-trainable legacy offsets no longer shrink view batches (#40).
|
|
15
|
+
- Compose shared legacy offsets with one launch per job: the frames of all views are staged in one transfer and `compose_surface` runs on a 2-D grid whose rows are views. The offset snapshot is staged once per call and pixel tile; offsets supplied by the caller, such as in the geometry VJP replay, are staged on every call (#40).
|
|
16
|
+
- Count the batched frame staging in the streamed surface memory budget, and pack shared offsets into the workspace without a device staging buffer (#40).
|
|
17
|
+
|
|
10
18
|
## [2.4.1] - 2026-10-20
|
|
11
19
|
|
|
12
20
|
### 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.4.
|
|
17
|
-
date-released: "2026-10-
|
|
16
|
+
version: 2.4.2
|
|
17
|
+
date-released: "2026-10-21"
|
|
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.4.
|
|
3
|
+
Version: 2.4.2
|
|
4
4
|
Summary: Differentiable CT projection and reconstruction: precision modes, streaming, distributed execution and geometry gradients
|
|
5
5
|
Project-URL: Homepage, https://github.com/sypsyp97/diffct
|
|
6
6
|
Project-URL: Documentation, https://sypsyp97.github.io/diffct/
|
|
@@ -51,11 +51,17 @@ def compose_surface(frames, offsets, source, points, view_index, pixel_start,
|
|
|
51
51
|
``frame32`` gives the arithmetic dtype of each packed frame component.
|
|
52
52
|
``offset32`` gives the offset arithmetic dtype. Values are stored exactly
|
|
53
53
|
in float64 staging buffers. Results use caller-owned float32 buffers.
|
|
54
|
+
|
|
55
|
+
A launch may compose several views with the same offsets: grid row ``y``
|
|
56
|
+
writes view ``view_index + y`` from the ``y``-th of ``gridDim.y`` equal
|
|
57
|
+
frame packets in ``frames``.
|
|
54
58
|
"""
|
|
55
59
|
pixel = cuda.grid(1)
|
|
60
|
+
view_index += cuda.blockIdx.y
|
|
61
|
+
base = cuda.blockIdx.y * (frames.size // cuda.gridDim.y)
|
|
56
62
|
if pixel == 0:
|
|
57
63
|
for axis in range(rank):
|
|
58
|
-
value = frames[axis]
|
|
64
|
+
value = frames[base + axis]
|
|
59
65
|
if not parallel:
|
|
60
66
|
value = _surface_add(value, -centre[axis], frame32[0])
|
|
61
67
|
source[view_index, axis] = np.float32(value)
|
|
@@ -85,20 +91,20 @@ def compose_surface(frames, offsets, source, points, view_index, pixel_start,
|
|
|
85
91
|
point_squared = np.float64(0.0)
|
|
86
92
|
for axis in range(rank):
|
|
87
93
|
point = _surface_add(
|
|
88
|
-
frames[rank + axis],
|
|
89
|
-
_surface_multiply(tangent, frames[2 * rank + axis], tangent32),
|
|
94
|
+
frames[base + rank + axis],
|
|
95
|
+
_surface_multiply(tangent, frames[base + 2 * rank + axis], tangent32),
|
|
90
96
|
frame32[1] and tangent32,
|
|
91
97
|
)
|
|
92
98
|
if rank == 3:
|
|
93
99
|
point = _surface_add(
|
|
94
|
-
point, _surface_multiply(vertical, frames[3 * rank + axis], vertical32),
|
|
100
|
+
point, _surface_multiply(vertical, frames[base + 3 * rank + axis], vertical32),
|
|
95
101
|
point32,
|
|
96
102
|
)
|
|
97
|
-
normal = frames[4 * rank + axis]
|
|
103
|
+
normal = frames[base + 4 * rank + axis]
|
|
98
104
|
else:
|
|
99
|
-
normal = frames[2 * rank + 1] if axis == 0 else -frames[2 * rank]
|
|
105
|
+
normal = frames[base + 2 * rank + 1] if axis == 0 else -frames[base + 2 * rank]
|
|
100
106
|
point = _surface_add(point, _surface_multiply(depth, normal, depth32), final32)
|
|
101
|
-
source_value = frames[axis]
|
|
107
|
+
source_value = frames[base + axis]
|
|
102
108
|
source_native = np.float32(source_value)
|
|
103
109
|
point_native = np.float32(point)
|
|
104
110
|
finite = (finite and math.isfinite(source_value) and math.isfinite(point)
|
|
@@ -489,6 +489,8 @@ class Projector:
|
|
|
489
489
|
if self.backend != "torch" and detector_surface is not None:
|
|
490
490
|
geometry, spec, surface_trainable, surface_metadata = self._surface_validation_inputs()
|
|
491
491
|
devices = self._devices or (torch.device("cuda", torch.cuda.current_device()),)
|
|
492
|
+
# Validation budgets depend on whether geometry gradients are needed.
|
|
493
|
+
self._geometry_trainable = bool(self._learnable) or surface_trainable
|
|
492
494
|
self._validate_surface_geometry(geometry, spec, devices)
|
|
493
495
|
except (TypeError, ValueError, RuntimeError) as error:
|
|
494
496
|
validation_error = error
|
|
@@ -769,6 +771,8 @@ class Projector:
|
|
|
769
771
|
|
|
770
772
|
def fits(shape):
|
|
771
773
|
points = surface or pixels != self.detector_shape
|
|
774
|
+
# Legacy offsets are one host snapshot; world points expand on the device.
|
|
775
|
+
host_points = points and self._host_world_points(geometry_spec)
|
|
772
776
|
buffers = _pipeline_buffers(self.beam, shape, batch, pixels, points,
|
|
773
777
|
geometry_grad=self._geometry_trainable,
|
|
774
778
|
ray_accumulators=extra_rays, peers=peers or foreign_output)
|
|
@@ -778,7 +782,7 @@ class Projector:
|
|
|
778
782
|
and _reservation_bytes(buffers) <= reservation_budget
|
|
779
783
|
and _largest_buffer_bytes(self.beam, shape, batch, pixels, points) <= allocation_limit
|
|
780
784
|
and max(self._surface_work_buffers(batch, pixels, spec=geometry_spec, streamed=True), default=0) <= allocation_limit
|
|
781
|
-
and _host_geometry_bytes(self.beam, batch, pixels,
|
|
785
|
+
and _host_geometry_bytes(self.beam, batch, pixels, host_points) <= _pipeline._GEOMETRY_WORK_BYTES
|
|
782
786
|
)
|
|
783
787
|
|
|
784
788
|
while not fits(minimum_shape):
|
|
@@ -826,7 +830,11 @@ class Projector:
|
|
|
826
830
|
rank, frames = len(self.volume_shape), 4 if self.beam == "cone" else 3
|
|
827
831
|
capacity = self._surface_capacity(views, pixels)
|
|
828
832
|
columns = 3 if self.beam == "cone" else 2
|
|
829
|
-
|
|
833
|
+
shared = spec is not None and spec.kind == "offsets" and spec.shared
|
|
834
|
+
if shared:
|
|
835
|
+
capacity = math.prod(pixels)
|
|
836
|
+
# Shared offsets pack the frame packets of every view of a job.
|
|
837
|
+
buffers = (_allocation_bytes(2 * (views if shared else 1) * (frames + (3 if rank == 3 else 0)) * rank),
|
|
830
838
|
_allocation_bytes(2 * (columns * capacity + 1)))
|
|
831
839
|
if self._geometry_trainable:
|
|
832
840
|
# Replay one view and a bounded detector row at a time.
|
|
@@ -850,7 +858,17 @@ class Projector:
|
|
|
850
858
|
dtype = torch.float32
|
|
851
859
|
for component in (*frames[1:], offsets):
|
|
852
860
|
dtype = torch.promote_types(dtype, component.dtype)
|
|
853
|
-
return _GeometrySpec("offsets", dtype=dtype)
|
|
861
|
+
return _GeometrySpec("offsets", dtype=dtype, shared=offsets.ndim == len(self.detector_shape) + 1)
|
|
862
|
+
|
|
863
|
+
def _host_world_points(self, spec):
|
|
864
|
+
"""Whether the host geometry budget must count world points of each view.
|
|
865
|
+
|
|
866
|
+
Shared legacy offsets are one host snapshot whose world points are
|
|
867
|
+
composed on the compute device; per-view offsets and geometry
|
|
868
|
+
gradients keep host work proportional to views x pixels.
|
|
869
|
+
"""
|
|
870
|
+
return not (spec is not None and spec.kind == "offsets" and spec.shared
|
|
871
|
+
and not self._geometry_trainable)
|
|
854
872
|
|
|
855
873
|
def _surface_validation_inputs(self):
|
|
856
874
|
"""Stage frame metadata once without expanding surface world points."""
|
|
@@ -884,7 +902,7 @@ class Projector:
|
|
|
884
902
|
buffers = self._surface_work_buffers(batch, pixels, spec=spec)
|
|
885
903
|
return (sum(buffers) <= budget and max(buffers, default=0) <= allocation
|
|
886
904
|
and _reservation_bytes(buffers) <= reservation
|
|
887
|
-
and _host_geometry_bytes(self.beam, batch, pixels,
|
|
905
|
+
and _host_geometry_bytes(self.beam, batch, pixels, self._host_world_points(spec))
|
|
888
906
|
<= _pipeline._GEOMETRY_WORK_BYTES)
|
|
889
907
|
|
|
890
908
|
if not fits():
|
|
@@ -1197,7 +1215,8 @@ class Projector:
|
|
|
1197
1215
|
estimate = size(shape, slots, retained, extra_rays)
|
|
1198
1216
|
reservation = (reserved(shape, slots, retained, extra_rays)
|
|
1199
1217
|
if estimate <= plan.gpu_budget_bytes else (1 << 63) - 1)
|
|
1200
|
-
|
|
1218
|
+
host_points = points and self._host_world_points(spec)
|
|
1219
|
+
host_fits = 2 * _host_geometry_bytes(self.beam, plan.view_chunk_size, pixels, host_points) <= _pipeline._GEOMETRY_WORK_BYTES
|
|
1201
1220
|
if (estimate > plan.gpu_budget_bytes
|
|
1202
1221
|
or reservation > plan.reservation_budget_bytes or not host_fits):
|
|
1203
1222
|
if schedule == "window":
|
|
@@ -1310,6 +1329,9 @@ class Projector:
|
|
|
1310
1329
|
rank, frame_count = len(self.volume_shape), 4 if self.beam == "cone" else 3
|
|
1311
1330
|
count = views.stop - views.start
|
|
1312
1331
|
pixel_shape = tuple(selected.stop - selected.start for selected in pixels)
|
|
1332
|
+
# Offsets sampled here derive from the snapshot and the pixel tile;
|
|
1333
|
+
# offsets supplied by the caller may differ per batch and are not cached.
|
|
1334
|
+
cacheable = offsets is None
|
|
1313
1335
|
if offsets is None:
|
|
1314
1336
|
offsets = _sample_offsets(self.beam, self.detector_shape, self.detector_spacing,
|
|
1315
1337
|
geometry, spec, views, pixels)
|
|
@@ -1324,15 +1346,42 @@ class Projector:
|
|
|
1324
1346
|
frame32 += (True,)
|
|
1325
1347
|
offset32 = ((torch.promote_types(torch.float32, offsets.dtype)
|
|
1326
1348
|
if spec.kind == "offsets" else common) == torch.float32)
|
|
1349
|
+
# Offsets shared by all views are staged once per call and composed
|
|
1350
|
+
# with one launch for all views instead of per view and row segment.
|
|
1351
|
+
shared = spec.kind == "offsets" and offsets.ndim == len(pixel_shape) + 1
|
|
1352
|
+
pixel_count = math.prod(pixel_shape)
|
|
1353
|
+
# One packed frame packet per view: source, center, axes, normal and
|
|
1354
|
+
# float32 cross-product scratch. Shared offsets pack every view of a job.
|
|
1355
|
+
stride = (frame_count + (3 if rank == 3 else 0)) * rank
|
|
1327
1356
|
if "_surface_workspace" not in groups:
|
|
1328
|
-
capacity = self._surface_capacity(count, pixel_shape)
|
|
1357
|
+
capacity = pixel_count if shared else self._surface_capacity(count, pixel_shape)
|
|
1329
1358
|
groups["_surface_workspace"] = (
|
|
1330
|
-
pipe.empty((
|
|
1359
|
+
pipe.empty((count if shared else 1) * stride, dtype=torch.float64),
|
|
1331
1360
|
pipe.empty((3 if rank == 3 else 2) * capacity + 1, dtype=torch.float64), capacity,
|
|
1332
1361
|
)
|
|
1333
1362
|
packed_frames, packed_offsets, capacity = groups["_surface_workspace"]
|
|
1334
1363
|
columns = (0, 1, 2) if rank == 3 else (0, 2)
|
|
1335
1364
|
source, points = targets[0].view(count, rank), targets[1].view(count, -1, rank)
|
|
1365
|
+
# A workspace sized for an earlier, smaller tile keeps the segment path.
|
|
1366
|
+
shared = shared and capacity >= pixel_count
|
|
1367
|
+
if shared and packed_frames.numel() < count * stride:
|
|
1368
|
+
# Earlier launches on the upload stream may still read the smaller buffer.
|
|
1369
|
+
pipe.retain(pipe.event(pipe.upload), packed_frames, None)
|
|
1370
|
+
packed_frames = pipe.empty(count * stride, dtype=torch.float64)
|
|
1371
|
+
groups["_surface_workspace"] = (packed_frames, packed_offsets, capacity)
|
|
1372
|
+
# Every job of a call slices the same offset snapshot: stage it once per
|
|
1373
|
+
# call, pixel tile and workspace. The key holds the snapshot and the
|
|
1374
|
+
# workspace, so neither a later snapshot reusing storage nor a recreated
|
|
1375
|
+
# workspace can match it.
|
|
1376
|
+
if shared and cacheable:
|
|
1377
|
+
staged_key = (geometry[frame_count], packed_offsets, tuple((part.start, part.stop) for part in pixels))
|
|
1378
|
+
restage = not ("_surface_shared_staged" in groups
|
|
1379
|
+
and groups["_surface_shared_staged"][0] is staged_key[0]
|
|
1380
|
+
and groups["_surface_shared_staged"][1] is staged_key[1]
|
|
1381
|
+
and groups["_surface_shared_staged"][2] == staged_key[2])
|
|
1382
|
+
else:
|
|
1383
|
+
restage = shared
|
|
1384
|
+
groups.pop("_surface_shared_staged", None)
|
|
1336
1385
|
with torch.cuda.device(pipe.device), torch.cuda.stream(pipe.upload), cuda.gpus[pipe.device.index]:
|
|
1337
1386
|
if origin is not None:
|
|
1338
1387
|
pipe.upload.wait_stream(origin)
|
|
@@ -1340,7 +1389,51 @@ class Projector:
|
|
|
1340
1389
|
stream = _get_numba_external_stream_for(pipe.upload)
|
|
1341
1390
|
arrays = tuple(TorchCUDABridge.tensor_to_cuda_array(value)
|
|
1342
1391
|
for value in (packed_frames, packed_offsets, source, points))
|
|
1343
|
-
|
|
1392
|
+
if restage:
|
|
1393
|
+
# Pack the float64 columns where the offsets live, then copy each
|
|
1394
|
+
# column into the workspace without a device staging buffer.
|
|
1395
|
+
flat = offsets.detach().reshape(pixel_count, 3)
|
|
1396
|
+
if offsets.is_cuda:
|
|
1397
|
+
for column, original in enumerate(columns):
|
|
1398
|
+
packed_offsets[column * capacity:column * capacity + pixel_count].copy_(flat[:, original])
|
|
1399
|
+
else:
|
|
1400
|
+
host = flat[:, list(columns)].T.to(dtype=torch.float64).contiguous()
|
|
1401
|
+
for column in range(len(columns)):
|
|
1402
|
+
_pipeline._copy(packed_offsets[column * capacity:column * capacity + pixel_count],
|
|
1403
|
+
host[column], pipe.stats)
|
|
1404
|
+
pipe.retain(pipe.event(pipe.upload), host, packed_offsets)
|
|
1405
|
+
if cacheable:
|
|
1406
|
+
groups["_surface_shared_staged"] = staged_key
|
|
1407
|
+
if shared:
|
|
1408
|
+
# Stage the frame packets of every view in one transfer and
|
|
1409
|
+
# compose all views with one launch; grid row y composes view y.
|
|
1410
|
+
block = packed_frames[:count * stride].view(count, stride)
|
|
1411
|
+
packet = torch.zeros((count, stride), dtype=torch.float64)
|
|
1412
|
+
for index, value in enumerate(frames):
|
|
1413
|
+
packet[:, index * rank:(index + 1) * rank].copy_(value[views].detach())
|
|
1414
|
+
_pipeline._copy(block, packet, pipe.stats)
|
|
1415
|
+
pipe.retain(pipe.event(pipe.upload), packet, packed_frames)
|
|
1416
|
+
if rank == 3:
|
|
1417
|
+
normal = block[:, 4 * rank:5 * rank]
|
|
1418
|
+
axis_u, axis_v = block[:, 2 * rank:3 * rank], block[:, 3 * rank:4 * rank]
|
|
1419
|
+
if frame32[2] and frame32[3]:
|
|
1420
|
+
scratch = block.view(torch.float32)[:, 10 * rank:13 * rank]
|
|
1421
|
+
scratch[:, :rank].copy_(axis_u)
|
|
1422
|
+
scratch[:, rank:2 * rank].copy_(axis_v)
|
|
1423
|
+
torch.linalg.cross(scratch[:, :rank], scratch[:, rank:2 * rank], dim=1,
|
|
1424
|
+
out=scratch[:, 2 * rank:])
|
|
1425
|
+
normal.copy_(scratch[:, 2 * rank:])
|
|
1426
|
+
else:
|
|
1427
|
+
torch.linalg.cross(axis_u, axis_v, dim=1, out=normal)
|
|
1428
|
+
for first in range(0, count, 65535):
|
|
1429
|
+
rows = min(65535, count - first)
|
|
1430
|
+
compose_surface[((pixel_count + 127) // 128, rows), 128, stream](
|
|
1431
|
+
TorchCUDABridge.tensor_to_cuda_array(block[first:first + rows].reshape(-1)),
|
|
1432
|
+
*arrays[1:], first, 0, pixel_count, capacity, rank,
|
|
1433
|
+
self.beam == "parallel", frame32, offset32,
|
|
1434
|
+
tuple(centre) + (0.,) * (3 - rank), self.voxel_spacing,
|
|
1435
|
+
)
|
|
1436
|
+
for local in range(0 if shared else count):
|
|
1344
1437
|
packet = torch.cat([value[views.start + local].detach().to(device="cpu", dtype=torch.float64)
|
|
1345
1438
|
for value in frames])
|
|
1346
1439
|
_pipeline._copy(packed_frames[:frame_count * rank], packet, pipe.stats)
|
|
@@ -2139,6 +2232,7 @@ class Projector:
|
|
|
2139
2232
|
slot.finish()
|
|
2140
2233
|
pipe.reclaim_output()
|
|
2141
2234
|
del local_groups["_surface_workspace"]
|
|
2235
|
+
local_groups.pop("_surface_shared_staged", None)
|
|
2142
2236
|
reduced = surface_batch_vjp(
|
|
2143
2237
|
self.beam, self.detector_shape, self.detector_spacing, inputs, offsets,
|
|
2144
2238
|
global_slice, pixels, self.voxel_spacing, pipe.device, parts,
|
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
# diffct
|
|
1
|
+
# diffct introduction
|
|
2
2
|
|
|
3
3
|
The film preserves the upstream dark, serif presentation: LaTeX typography,
|
|
4
4
|
numbered chapters, a thin progress rule, framed CT images and monospace code
|
|
@@ -89,7 +89,7 @@ reconstruction panels retain their recorded inputs.
|
|
|
89
89
|
|
|
90
90
|
| Shot | Content checked | Evidence |
|
|
91
91
|
|---|---|---|
|
|
92
|
-
| Opening |
|
|
92
|
+
| Opening | Title, five chapters, measured CT surface | `walnut_surface.npz` |
|
|
93
93
|
| Trajectories | Five schematic paths; source/centre/u/v tuple; valid per-view poses | [Projector](../../diffct/operators.py), unchanged [ScanDiagram](scan_diagram.py) |
|
|
94
94
|
| Autograd | Simulated parallel sinogram; Siddon path sum; matched adjoint; ordinary squared-loss gradient; batch/second data derivatives | [Projector](../../diffct/operators.py), [portable traversal](../../diffct/backends.py), `make_inputs.py`, `data2d.npz` |
|
|
95
95
|
| Geometry | First-order source/centre/axis gradients; schematic pose dots; log10 loss, 150 steps, angle RMSE 0.55° to 0.006° | [geometry VJP](../../diffct/kernels/geometry_vjp.py), `calib_history.json` |
|
|
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "diffct"
|
|
7
|
-
version = "2.4.
|
|
7
|
+
version = "2.4.2"
|
|
8
8
|
description = "Differentiable CT projection and reconstruction: precision modes, streaming, distributed execution and geometry gradients"
|
|
9
9
|
readme = "README.md"
|
|
10
10
|
requires-python = ">=3.10"
|
|
@@ -85,11 +85,13 @@ class _KernelEntries:
|
|
|
85
85
|
|
|
86
86
|
def __getitem__(self, configuration):
|
|
87
87
|
launch = self.kernel[configuration]
|
|
88
|
+
grid = configuration[0]
|
|
89
|
+
rows = grid[1] if isinstance(grid, tuple) and len(grid) > 1 else 1
|
|
88
90
|
|
|
89
91
|
def run(*args):
|
|
90
92
|
frames, offsets, source, points, view, start, count, capacity = args[:8]
|
|
91
|
-
self.records.append(dict(frames=frames.copy(), source=source, points=points,
|
|
92
|
-
view=view, start=start, count=count, capacity=capacity))
|
|
93
|
+
self.records.append(dict(frames=frames.copy(), frame_storage=frames, source=source, points=points,
|
|
94
|
+
view=view, rows=rows, start=start, count=count, capacity=capacity))
|
|
93
95
|
return launch(*args)
|
|
94
96
|
|
|
95
97
|
return run
|
|
@@ -177,6 +179,8 @@ def test_fused_helper_runs_real_kernel_for_global_batches_into_reused_bounded_ta
|
|
|
177
179
|
first_pixels = (slice(2, 7), slice(3, 8)) if beam == "cone" else (slice(2, 27),)
|
|
178
180
|
tail_pixels = (slice(8, 11), slice(8, 13)) if beam == "cone" else (slice(26, 29),)
|
|
179
181
|
batches = [(slice(2, 4), first_pixels), (slice(6, 7), tail_pixels)]
|
|
182
|
+
shared_offsets = kind == "offsets" and not per_view
|
|
183
|
+
packet_size = (len(case.trajectory) + (3 if rank == 3 else 0)) * rank
|
|
180
184
|
centre = (.1875, -.3125, .4375)[:rank]
|
|
181
185
|
source_buffer, point_buffer = torch.empty(2 * rank), torch.empty(2 * 25 * rank)
|
|
182
186
|
original_storage = (source_buffer.data_ptr(), point_buffer.data_ptr())
|
|
@@ -213,26 +217,59 @@ def test_fused_helper_runs_real_kernel_for_global_batches_into_reused_bounded_ta
|
|
|
213
217
|
assert torch.all(point_buffer[math.prod(shape):] == 77.)
|
|
214
218
|
assert (source_buffer.data_ptr(), point_buffer.data_ptr()) == original_storage
|
|
215
219
|
observed = entries.records[before:]
|
|
216
|
-
assert observed
|
|
220
|
+
assert observed
|
|
221
|
+
pixel_count = math.prod(shape[1:-1])
|
|
222
|
+
coverage = np.zeros((count, pixel_count), dtype=np.int64)
|
|
223
|
+
if shared_offsets:
|
|
224
|
+
assert len(observed) == math.ceil(count / 65535)
|
|
225
|
+
assert sum(record["rows"] for record in observed) == count
|
|
217
226
|
for record in observed:
|
|
227
|
+
assert np.shares_memory(record["frame_storage"], groups["_surface_workspace"][0].numpy())
|
|
218
228
|
assert np.shares_memory(record["source"], source_buffer.numpy())
|
|
219
229
|
assert np.shares_memory(record["points"], point_buffer.numpy())
|
|
230
|
+
assert 0 <= record["view"] < record["view"] + record["rows"] <= count
|
|
231
|
+
assert 1 <= record["rows"] <= 65535
|
|
220
232
|
assert record["count"] <= record["capacity"]
|
|
221
|
-
assert record["start"] + record["count"] <=
|
|
222
|
-
|
|
223
|
-
|
|
224
|
-
|
|
225
|
-
|
|
226
|
-
|
|
227
|
-
|
|
228
|
-
|
|
229
|
-
|
|
230
|
-
|
|
231
|
-
|
|
232
|
-
|
|
233
|
+
assert 0 <= record["start"] < record["start"] + record["count"] <= pixel_count
|
|
234
|
+
if shared_offsets:
|
|
235
|
+
assert record["start"] == 0 and record["count"] == pixel_count
|
|
236
|
+
else:
|
|
237
|
+
assert record["rows"] == 1
|
|
238
|
+
coverage[record["view"]:record["view"] + record["rows"],
|
|
239
|
+
record["start"]:record["start"] + record["count"]] += 1
|
|
240
|
+
packets = record["frames"].reshape(record["rows"], packet_size)
|
|
241
|
+
assert record["frames"].nbytes <= record["rows"] * packet_size * 8
|
|
242
|
+
for row, packet in enumerate(packets):
|
|
243
|
+
frame_values = [value[views.start + record["view"] + row].numpy()
|
|
244
|
+
for value in case.trajectory]
|
|
245
|
+
original_prefix = np.concatenate([value.astype(np.float64) for value in frame_values])
|
|
246
|
+
np.testing.assert_array_equal(packet[:original_prefix.size], original_prefix)
|
|
247
|
+
if rank == 3:
|
|
248
|
+
normal_dtype = np.result_type(np.float32, frame_values[2].dtype, frame_values[3].dtype)
|
|
249
|
+
if kind == "sampled":
|
|
250
|
+
normal_dtype = np.result_type(normal_dtype, *(value.dtype for value in frame_values), local_offsets.dtype)
|
|
251
|
+
normal = np.cross(frame_values[2].astype(normal_dtype), frame_values[3].astype(normal_dtype))
|
|
252
|
+
np.testing.assert_array_equal(packet[12:15], normal.astype(np.float64))
|
|
253
|
+
np.testing.assert_array_equal(coverage, np.ones_like(coverage))
|
|
254
|
+
assert len(pipe.allocations) == 2, \
|
|
255
|
+
"packed workspaces were allocated again for the tail batch"
|
|
256
|
+
packed_frames, packed_offsets, capacity = groups["_surface_workspace"]
|
|
257
|
+
assert "_surface_frames" not in groups
|
|
258
|
+
assert packed_frames is pipe.allocations[0] and packed_offsets is pipe.allocations[1]
|
|
259
|
+
maximum_views = max(views.stop - views.start for views, _ in batches)
|
|
260
|
+
assert packed_frames.dtype == torch.float64
|
|
261
|
+
assert packed_frames.numel() == (maximum_views if shared_offsets else 1) * packet_size
|
|
233
262
|
workspace_bytes = sum(math.ceil(value.numel() * value.element_size() / 512) * 512 for value in pipe.allocations)
|
|
234
263
|
target_bytes = sum(math.ceil(value.numel() * value.element_size() / 512) * 512 for value in (source_buffer, point_buffer))
|
|
235
|
-
|
|
264
|
+
maximum_bytes = target_bytes
|
|
265
|
+
if shared_offsets:
|
|
266
|
+
maximum_pixels = max(math.prod(part.stop - part.start for part in pixels) for _, pixels in batches)
|
|
267
|
+
assert capacity == maximum_pixels
|
|
268
|
+
# One packet per view and one shared offset tile plus status.
|
|
269
|
+
maximum_bytes = sum(math.ceil(size / 512) * 512 for size in (
|
|
270
|
+
maximum_views * packet_size * 8, (rank * maximum_pixels + 1) * 8,
|
|
271
|
+
))
|
|
272
|
+
assert workspace_bytes <= maximum_bytes, "packed staging exceeded the geometry workspace budget"
|
|
236
273
|
assert any(record["count"] < record["capacity"] for record in entries.records)
|
|
237
274
|
assert len(samples) == (2 if kind == "sampled" else 0), "surface sampler ran more than once per native batch"
|
|
238
275
|
if kind == "sampled":
|
|
@@ -246,6 +283,60 @@ def test_fused_helper_runs_real_kernel_for_global_batches_into_reused_bounded_ta
|
|
|
246
283
|
assert bool(pipe.waits) == (offset_origin == "cuda-metadata")
|
|
247
284
|
|
|
248
285
|
|
|
286
|
+
@pytest.mark.parametrize("beam", ["cone", "fan", "parallel"])
|
|
287
|
+
@pytest.mark.parametrize("first_count", [1, 3], ids=["grow-frames", "reuse-frames"])
|
|
288
|
+
def test_fused_helper_restages_caller_supplied_shared_offsets_for_geometry_vjp(beam, first_count):
|
|
289
|
+
case = _case(beam, views=5)
|
|
290
|
+
projector = _projector(case, backend="torch")
|
|
291
|
+
rank = len(case.shape)
|
|
292
|
+
pixels = (slice(1, 3), slice(2, 5)) if rank == 3 else (slice(1, 5),)
|
|
293
|
+
pixel_shape = tuple(part.stop - part.start for part in pixels)
|
|
294
|
+
count = 3
|
|
295
|
+
full_offsets = torch.zeros((*case.detector, 3), dtype=torch.float64)
|
|
296
|
+
geometry, spec = (*case.trajectory, full_offsets), _GeometrySpec("offsets")
|
|
297
|
+
first_offsets = torch.zeros((*pixel_shape, 3), dtype=torch.float64, requires_grad=True)
|
|
298
|
+
second_offsets = torch.tensor([.375, -.25, .625], dtype=torch.float64).expand(
|
|
299
|
+
*pixel_shape, 3).clone().requires_grad_()
|
|
300
|
+
targets = (torch.empty((count, rank)), torch.empty((count, *pixel_shape, rank)))
|
|
301
|
+
centre = (.1875, -.3125, .4375)[:rank]
|
|
302
|
+
groups, pipe = {}, _CpuGeometryPipe()
|
|
303
|
+
entries = _KernelEntries(importlib.import_module("diffct.kernels.surface_geometry").compose_surface)
|
|
304
|
+
outputs = []
|
|
305
|
+
with _cpu_geometry_transport(entries), \
|
|
306
|
+
mock.patch.object(operators, "_sample_offsets", side_effect=AssertionError("caller supplied offsets")):
|
|
307
|
+
for actual_count, offsets in ((first_count, first_offsets), (count, second_offsets)):
|
|
308
|
+
views = slice(1, 1 + actual_count)
|
|
309
|
+
selected_targets = tuple(value[:actual_count] for value in targets)
|
|
310
|
+
projector._surface_geometry_into(geometry, spec, views, pixels, centre,
|
|
311
|
+
selected_targets, pipe, groups, offsets=offsets)
|
|
312
|
+
frames = tuple(value[views].numpy().copy() for value in case.trajectory)
|
|
313
|
+
local_offsets = np.broadcast_to(offsets.detach().numpy(), (actual_count, *pixel_shape, 3))
|
|
314
|
+
expected = (_world(frames, local_offsets, beam) - np.asarray(centre)).astype(np.float32)
|
|
315
|
+
expected_source = frames[0] if beam == "parallel" else frames[0] - np.asarray(centre)
|
|
316
|
+
np.testing.assert_array_equal(selected_targets[0].numpy(), expected_source.astype(np.float32))
|
|
317
|
+
np.testing.assert_array_equal(selected_targets[1].numpy(), expected)
|
|
318
|
+
outputs.append(selected_targets[1].clone())
|
|
319
|
+
packet_size = (len(case.trajectory) + (3 if rank == 3 else 0)) * rank
|
|
320
|
+
packed_frames, packed_offsets, capacity = groups["_surface_workspace"]
|
|
321
|
+
assert "_surface_frames" not in groups
|
|
322
|
+
assert packed_frames.dtype == torch.float64
|
|
323
|
+
assert packed_frames.numel() == actual_count * packet_size
|
|
324
|
+
assert np.shares_memory(entries.records[-1]["frame_storage"], packed_frames.numpy())
|
|
325
|
+
if len(outputs) == 1:
|
|
326
|
+
workspace = groups["_surface_workspace"]
|
|
327
|
+
assert len(pipe.allocations) == 2
|
|
328
|
+
assert packed_frames is pipe.allocations[0] and packed_offsets is pipe.allocations[1]
|
|
329
|
+
else:
|
|
330
|
+
assert (groups["_surface_workspace"] is workspace) == (first_count == count)
|
|
331
|
+
assert (packed_frames is workspace[0]) == (first_count == count)
|
|
332
|
+
assert packed_offsets is workspace[1] and capacity == workspace[2]
|
|
333
|
+
assert len(pipe.allocations) == (3 if first_count < count else 2)
|
|
334
|
+
assert packed_frames is pipe.allocations[-1 if first_count < count else 0]
|
|
335
|
+
assert len(entries.records) == 2
|
|
336
|
+
assert [record["rows"] for record in entries.records] == [first_count, count]
|
|
337
|
+
assert not torch.equal(outputs[0], outputs[1][:first_count])
|
|
338
|
+
|
|
339
|
+
|
|
249
340
|
@pytest.mark.parametrize("failure,message", [
|
|
250
341
|
("nonfinite", "finite"), ("coincident", "must differ"),
|
|
251
342
|
("near", "1e6"), ("far", "1e15"),
|