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.
Files changed (146) hide show
  1. {diffct-2.4.1 → diffct-2.4.2}/CHANGELOG.md +8 -0
  2. {diffct-2.4.1 → diffct-2.4.2}/CITATION.cff +2 -2
  3. {diffct-2.4.1 → diffct-2.4.2}/PKG-INFO +1 -1
  4. {diffct-2.4.1 → diffct-2.4.2}/diffct/__init__.py +1 -1
  5. {diffct-2.4.1 → diffct-2.4.2}/diffct/kernels/surface_geometry.py +13 -7
  6. {diffct-2.4.1 → diffct-2.4.2}/diffct/operators.py +102 -8
  7. {diffct-2.4.1 → diffct-2.4.2}/diffct/surfaces.py +2 -0
  8. {diffct-2.4.1 → diffct-2.4.2}/docs/video/README.md +2 -2
  9. {diffct-2.4.1 → diffct-2.4.2}/pyproject.toml +1 -1
  10. {diffct-2.4.1 → diffct-2.4.2}/tests/test_surface_geometry_cudasim.py +107 -16
  11. diffct-2.4.2/tests/test_surface_host_sampling_budget.py +257 -0
  12. {diffct-2.4.1 → diffct-2.4.2}/tests/test_surface_pipeline_contracts.py +20 -0
  13. {diffct-2.4.1 → diffct-2.4.2}/.gitignore +0 -0
  14. {diffct-2.4.1 → diffct-2.4.2}/LICENSE +0 -0
  15. {diffct-2.4.1 → diffct-2.4.2}/NOTICE +0 -0
  16. {diffct-2.4.1 → diffct-2.4.2}/README.md +0 -0
  17. {diffct-2.4.1 → diffct-2.4.2}/diffct/_solver_state.py +0 -0
  18. {diffct-2.4.1 → diffct-2.4.2}/diffct/_solver_tv.py +0 -0
  19. {diffct-2.4.1 → diffct-2.4.2}/diffct/analytical.py +0 -0
  20. {diffct-2.4.1 → diffct-2.4.2}/diffct/analytical_execution.py +0 -0
  21. {diffct-2.4.1 → diffct-2.4.2}/diffct/backends.py +0 -0
  22. {diffct-2.4.1 → diffct-2.4.2}/diffct/chunking.py +0 -0
  23. {diffct-2.4.1 → diffct-2.4.2}/diffct/constants.py +0 -0
  24. {diffct-2.4.1 → diffct-2.4.2}/diffct/differentiable.py +0 -0
  25. {diffct-2.4.1 → diffct-2.4.2}/diffct/distribution.py +0 -0
  26. {diffct-2.4.1 → diffct-2.4.2}/diffct/geometry.py +0 -0
  27. {diffct-2.4.1 → diffct-2.4.2}/diffct/kernels/__init__.py +0 -0
  28. {diffct-2.4.1 → diffct-2.4.2}/diffct/kernels/analytical_transpose.py +0 -0
  29. {diffct-2.4.1 → diffct-2.4.2}/diffct/kernels/cone_beam.py +0 -0
  30. {diffct-2.4.1 → diffct-2.4.2}/diffct/kernels/fan_beam.py +0 -0
  31. {diffct-2.4.1 → diffct-2.4.2}/diffct/kernels/geometry_vjp.py +0 -0
  32. {diffct-2.4.1 → diffct-2.4.2}/diffct/kernels/parallel_beam.py +0 -0
  33. {diffct-2.4.1 → diffct-2.4.2}/diffct/pipeline.py +0 -0
  34. {diffct-2.4.1 → diffct-2.4.2}/diffct/projectors.py +0 -0
  35. {diffct-2.4.1 → diffct-2.4.2}/diffct/reconstruction.py +0 -0
  36. {diffct-2.4.1 → diffct-2.4.2}/diffct/solvers.py +0 -0
  37. {diffct-2.4.1 → diffct-2.4.2}/diffct/storage.py +0 -0
  38. {diffct-2.4.1 → diffct-2.4.2}/diffct/streaming.py +0 -0
  39. {diffct-2.4.1 → diffct-2.4.2}/diffct/surface_gradients.py +0 -0
  40. {diffct-2.4.1 → diffct-2.4.2}/diffct/utils.py +0 -0
  41. {diffct-2.4.1 → diffct-2.4.2}/examples/README.md +0 -0
  42. {diffct-2.4.1 → diffct-2.4.2}/examples/_block_cgls.py +0 -0
  43. {diffct-2.4.1 → diffct-2.4.2}/examples/_common.py +0 -0
  44. {diffct-2.4.1 → diffct-2.4.2}/examples/data/NOTICE +0 -0
  45. {diffct-2.4.1 → diffct-2.4.2}/examples/data/preprocess_walnut.py +0 -0
  46. {diffct-2.4.1 → diffct-2.4.2}/examples/walnut_reconstruction.py +0 -0
  47. {diffct-2.4.1 → diffct-2.4.2}/formal/README.md +0 -0
  48. {diffct-2.4.1 → diffct-2.4.2}/pytest.ini +0 -0
  49. {diffct-2.4.1 → diffct-2.4.2}/tests/__init__.py +0 -0
  50. {diffct-2.4.1 → diffct-2.4.2}/tests/_phantoms.py +0 -0
  51. {diffct-2.4.1 → diffct-2.4.2}/tests/benchmarks/README.md +0 -0
  52. {diffct-2.4.1 → diffct-2.4.2}/tests/benchmarks/__init__.py +0 -0
  53. {diffct-2.4.1 → diffct-2.4.2}/tests/benchmarks/_common.py +0 -0
  54. {diffct-2.4.1 → diffct-2.4.2}/tests/benchmarks/test_bench_cone.py +0 -0
  55. {diffct-2.4.1 → diffct-2.4.2}/tests/benchmarks/test_bench_fan.py +0 -0
  56. {diffct-2.4.1 → diffct-2.4.2}/tests/benchmarks/test_bench_parallel.py +0 -0
  57. {diffct-2.4.1 → diffct-2.4.2}/tests/distributed_projector_check.py +0 -0
  58. {diffct-2.4.1 → diffct-2.4.2}/tests/test_adjoint_inner_product.py +0 -0
  59. {diffct-2.4.1 → diffct-2.4.2}/tests/test_allocator_capacity.py +0 -0
  60. {diffct-2.4.1 → diffct-2.4.2}/tests/test_analytical_autograd.py +0 -0
  61. {diffct-2.4.1 → diffct-2.4.2}/tests/test_analytical_distribution.py +0 -0
  62. {diffct-2.4.1 → diffct-2.4.2}/tests/test_analytical_regressions.py +0 -0
  63. {diffct-2.4.1 → diffct-2.4.2}/tests/test_analytical_weight_distribution.py +0 -0
  64. {diffct-2.4.1 → diffct-2.4.2}/tests/test_angular_low_precision.py +0 -0
  65. {diffct-2.4.1 → diffct-2.4.2}/tests/test_axis_aligned_chords.py +0 -0
  66. {diffct-2.4.1 → diffct-2.4.2}/tests/test_batch_execution_edges.py +0 -0
  67. {diffct-2.4.1 → diffct-2.4.2}/tests/test_batched_surfaces.py +0 -0
  68. {diffct-2.4.1 → diffct-2.4.2}/tests/test_block_cgls.py +0 -0
  69. {diffct-2.4.1 → diffct-2.4.2}/tests/test_block_storage.py +0 -0
  70. {diffct-2.4.1 → diffct-2.4.2}/tests/test_call_trajectory.py +0 -0
  71. {diffct-2.4.1 → diffct-2.4.2}/tests/test_cgls_state_admission.py +0 -0
  72. {diffct-2.4.1 → diffct-2.4.2}/tests/test_checkpoint_metadata_validation.py +0 -0
  73. {diffct-2.4.1 → diffct-2.4.2}/tests/test_chunked_projector.py +0 -0
  74. {diffct-2.4.1 → diffct-2.4.2}/tests/test_completed_transfer_reuse.py +0 -0
  75. {diffct-2.4.1 → diffct-2.4.2}/tests/test_cone_backprojection_transpose.py +0 -0
  76. {diffct-2.4.1 → diffct-2.4.2}/tests/test_cone_projector_autograd.py +0 -0
  77. {diffct-2.4.1 → diffct-2.4.2}/tests/test_cpu_install.py +0 -0
  78. {diffct-2.4.1 → diffct-2.4.2}/tests/test_cuda_exit_deallocations.py +0 -0
  79. {diffct-2.4.1 → diffct-2.4.2}/tests/test_cuda_sink_lifetimes.py +0 -0
  80. {diffct-2.4.1 → diffct-2.4.2}/tests/test_cuda_smoke.py +0 -0
  81. {diffct-2.4.1 → diffct-2.4.2}/tests/test_detector_surface_dtypes.py +0 -0
  82. {diffct-2.4.1 → diffct-2.4.2}/tests/test_detector_surface_mapping.py +0 -0
  83. {diffct-2.4.1 → diffct-2.4.2}/tests/test_detector_surface_units.py +0 -0
  84. {diffct-2.4.1 → diffct-2.4.2}/tests/test_detector_surfaces.py +0 -0
  85. {diffct-2.4.1 → diffct-2.4.2}/tests/test_distributed_training_transforms.py +0 -0
  86. {diffct-2.4.1 → diffct-2.4.2}/tests/test_example_psnr.py +0 -0
  87. {diffct-2.4.1 → diffct-2.4.2}/tests/test_execution_buffers.py +0 -0
  88. {diffct-2.4.1 → diffct-2.4.2}/tests/test_execution_features.py +0 -0
  89. {diffct-2.4.1 → diffct-2.4.2}/tests/test_execution_schedules.py +0 -0
  90. {diffct-2.4.1 → diffct-2.4.2}/tests/test_fbp_fan_accuracy.py +0 -0
  91. {diffct-2.4.1 → diffct-2.4.2}/tests/test_fbp_fan_offsets.py +0 -0
  92. {diffct-2.4.1 → diffct-2.4.2}/tests/test_fbp_parallel_accuracy.py +0 -0
  93. {diffct-2.4.1 → diffct-2.4.2}/tests/test_fdk_cone_accuracy.py +0 -0
  94. {diffct-2.4.1 → diffct-2.4.2}/tests/test_fdk_cone_offsets.py +0 -0
  95. {diffct-2.4.1 → diffct-2.4.2}/tests/test_geometry_autograd.py +0 -0
  96. {diffct-2.4.1 → diffct-2.4.2}/tests/test_geometry_conversions.py +0 -0
  97. {diffct-2.4.1 → diffct-2.4.2}/tests/test_geometry_gradients.py +0 -0
  98. {diffct-2.4.1 → diffct-2.4.2}/tests/test_gradcheck.py +0 -0
  99. {diffct-2.4.1 → diffct-2.4.2}/tests/test_lazy_native_observers.py +0 -0
  100. {diffct-2.4.1 → diffct-2.4.2}/tests/test_live_trajectory_validation.py +0 -0
  101. {diffct-2.4.1 → diffct-2.4.2}/tests/test_multi_gpu_exit.py +0 -0
  102. {diffct-2.4.1 → diffct-2.4.2}/tests/test_native_endpoint_precision.py +0 -0
  103. {diffct-2.4.1 → diffct-2.4.2}/tests/test_native_surface_transforms.py +0 -0
  104. {diffct-2.4.1 → diffct-2.4.2}/tests/test_native_volume_layout.py +0 -0
  105. {diffct-2.4.1 → diffct-2.4.2}/tests/test_native_volume_layout_distributed.py +0 -0
  106. {diffct-2.4.1 → diffct-2.4.2}/tests/test_output_pipeline_overlap.py +0 -0
  107. {diffct-2.4.1 → diffct-2.4.2}/tests/test_parallel3d.py +0 -0
  108. {diffct-2.4.1 → diffct-2.4.2}/tests/test_parallel3d_surface_oracle.py +0 -0
  109. {diffct-2.4.1 → diffct-2.4.2}/tests/test_partition_lifecycle.py +0 -0
  110. {diffct-2.4.1 → diffct-2.4.2}/tests/test_plan_memory_bounds.py +0 -0
  111. {diffct-2.4.1 → diffct-2.4.2}/tests/test_portable_distributed_geometry.py +0 -0
  112. {diffct-2.4.1 → diffct-2.4.2}/tests/test_portable_endpoint_precision.py +0 -0
  113. {diffct-2.4.1 → diffct-2.4.2}/tests/test_portable_execution_bounds.py +0 -0
  114. {diffct-2.4.1 → diffct-2.4.2}/tests/test_portable_numeric_regressions.py +0 -0
  115. {diffct-2.4.1 → diffct-2.4.2}/tests/test_portable_sampler_snapshot.py +0 -0
  116. {diffct-2.4.1 → diffct-2.4.2}/tests/test_portable_surface_admission.py +0 -0
  117. {diffct-2.4.1 → diffct-2.4.2}/tests/test_projector_api.py +0 -0
  118. {diffct-2.4.1 → diffct-2.4.2}/tests/test_projector_degenerate_geometry.py +0 -0
  119. {diffct-2.4.1 → diffct-2.4.2}/tests/test_projector_geometry_devices.py +0 -0
  120. {diffct-2.4.1 → diffct-2.4.2}/tests/test_ramp_filter_windows.py +0 -0
  121. {diffct-2.4.1 → diffct-2.4.2}/tests/test_reconstruction_api.py +0 -0
  122. {diffct-2.4.1 → diffct-2.4.2}/tests/test_reference_backend.py +0 -0
  123. {diffct-2.4.1 → diffct-2.4.2}/tests/test_release_2_0_1.py +0 -0
  124. {diffct-2.4.1 → diffct-2.4.2}/tests/test_review_230.py +0 -0
  125. {diffct-2.4.1 → diffct-2.4.2}/tests/test_review_regressions.py +0 -0
  126. {diffct-2.4.1 → diffct-2.4.2}/tests/test_solver_admission_regressions.py +0 -0
  127. {diffct-2.4.1 → diffct-2.4.2}/tests/test_solver_api.py +0 -0
  128. {diffct-2.4.1 → diffct-2.4.2}/tests/test_solver_checkpoint_edges.py +0 -0
  129. {diffct-2.4.1 → diffct-2.4.2}/tests/test_solver_collective_failures.py +0 -0
  130. {diffct-2.4.1 → diffct-2.4.2}/tests/test_spatial_partition.py +0 -0
  131. {diffct-2.4.1 → diffct-2.4.2}/tests/test_surface_collective_failures.py +0 -0
  132. {diffct-2.4.1 → diffct-2.4.2}/tests/test_surface_composition_kernel.py +0 -0
  133. {diffct-2.4.1 → diffct-2.4.2}/tests/test_surface_compute_shards.py +0 -0
  134. {diffct-2.4.1 → diffct-2.4.2}/tests/test_surface_gpu_execution.py +0 -0
  135. {diffct-2.4.1 → diffct-2.4.2}/tests/test_surface_gradient_replay.py +0 -0
  136. {diffct-2.4.1 → diffct-2.4.2}/tests/test_surface_native_regressions.py +0 -0
  137. {diffct-2.4.1 → diffct-2.4.2}/tests/test_surface_parameter_updates.py +0 -0
  138. {diffct-2.4.1 → diffct-2.4.2}/tests/test_surface_trajectory_transform.py +0 -0
  139. {diffct-2.4.1 → diffct-2.4.2}/tests/test_training_transforms.py +0 -0
  140. {diffct-2.4.1 → diffct-2.4.2}/tests/test_trajectory_admission_cost.py +0 -0
  141. {diffct-2.4.1 → diffct-2.4.2}/tests/test_trajectory_surface_integration.py +0 -0
  142. {diffct-2.4.1 → diffct-2.4.2}/tests/test_transfer_pipeline.py +0 -0
  143. {diffct-2.4.1 → diffct-2.4.2}/tests/test_transfer_retirement_lifetimes.py +0 -0
  144. {diffct-2.4.1 → diffct-2.4.2}/tests/test_tv_zero_dependency.py +0 -0
  145. {diffct-2.4.1 → diffct-2.4.2}/tests/test_two_far_endpoints.py +0 -0
  146. {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.1
17
- date-released: "2026-10-20"
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.1
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/
@@ -64,7 +64,7 @@ try:
64
64
  except ImportError:
65
65
  pass
66
66
 
67
- __version__ = '2.4.1'
67
+ __version__ = '2.4.2'
68
68
 
69
69
  __all__ = [
70
70
  'Projector',
@@ -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, points) <= _pipeline._GEOMETRY_WORK_BYTES
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
- buffers = (_allocation_bytes(2 * (frames + (3 if rank == 3 else 0)) * rank),
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, True)
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
- host_fits = 2 * _host_geometry_bytes(self.beam, plan.view_chunk_size, pixels, points) <= _pipeline._GEOMETRY_WORK_BYTES
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((frame_count + (3 if rank == 3 else 0)) * rank, dtype=torch.float64),
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
- for local in range(count):
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,
@@ -10,6 +10,8 @@ class _GeometrySpec:
10
10
  kind: str
11
11
  sampler: object = None
12
12
  dtype: object = None
13
+ # One (U, V, 3) host offset snapshot shared by all views.
14
+ shared: bool = False
13
15
 
14
16
 
15
17
  class ParameterizedSurface:
@@ -1,4 +1,4 @@
1
- # diffct 2.3.0 introduction
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 | 2.3.0 title, five chapters, measured CT surface | `pyproject.toml`, `walnut_surface.npz` |
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.1"
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 and {record["view"] for record in observed} == set(range(count))
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"] <= math.prod(shape[1:-1])
222
- frame_values = [value[views.start + record["view"]].numpy() for value in case.trajectory]
223
- original_prefix = np.concatenate([value.astype(np.float64) for value in frame_values])
224
- np.testing.assert_array_equal(record["frames"][:original_prefix.size], original_prefix)
225
- assert record["frames"].nbytes <= 512, "packed frame arena exceeded one native allocation unit"
226
- if rank == 3:
227
- normal_dtype = np.result_type(np.float32, frame_values[2].dtype, frame_values[3].dtype)
228
- if kind == "sampled":
229
- normal_dtype = np.result_type(normal_dtype, *(value.dtype for value in frame_values), local_offsets.dtype)
230
- normal = np.cross(frame_values[2].astype(normal_dtype), frame_values[3].astype(normal_dtype))
231
- np.testing.assert_array_equal(record["frames"][12:15], normal.astype(np.float64))
232
- assert len(pipe.allocations) == 2, "packed workspaces were allocated again for the tail batch"
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
- assert workspace_bytes <= target_bytes, "packed staging exceeded the native geometry buffer bytes"
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"),