Skip to content

Add GPU acceleration subpackage linumpy.gpu (CuPy + CPU fallback) - #140

Open
FIrgolitsch wants to merge 1 commit into
linum-uqam:mainfrom
FIrgolitsch:pr-gpu-core
Open

Add GPU acceleration subpackage linumpy.gpu (CuPy + CPU fallback)#140
FIrgolitsch wants to merge 1 commit into
linum-uqam:mainfrom
FIrgolitsch:pr-gpu-core

Conversation

@FIrgolitsch

@FIrgolitsch FIrgolitsch commented Jul 7, 2026

Copy link
Copy Markdown
Contributor

Summary

Adds a new linumpy.gpu subpackage: GPU-accelerated implementations of the compute-heavy operations used elsewhere in the library (FFT phase correlation, affine interpolation, morphology, B-spline fitting, bias-field correction, image-quality metrics, tissue-interface detection, and Zarr I/O), all with automatic fallback to NumPy when CuPy/CUDA is unavailable.

Land after #139.

What changed

New package linumpy/gpu/ (13 modules, all new):

Module Provides
__init__.py Device detection, multi-GPU selection (file-locked round-robin across physical devices, honouring CUDA_VISIBLE_DEVICES), to_gpu/to_cpu/get_array_module helpers, list_gpus/select_gpu/select_best_gpu, env-gated via LINUMPY_USE_GPU
interface.py GPU tissue-interface detection (mirrors linumpy.geometry.interface)
array_ops.py GPU array utilities
fft_ops.py GPU phase correlation / FFT ops
interpolation.py GPU affine transform + resampling
morphology.py GPU morphological filters
corrections.py GPU correction kernels
bspline.py GPU B-spline fit/evaluate
bias_field.py GPU bias-field correction
image_quality.py GPU image-quality metrics
zarr_io.py GPU-aware Zarr I/O
kvikio_zarr.py GPUDirect-Storage (kvikio) Zarr reader
nvcomp_zstd.py GPU ZSTD compression via nvCOMP

Plus linumpy/tests/test_gpu_bspline.py — B-spline correctness (partition-of-unity, constant/linear recovery, mask handling) and CPU/GPU numerical agreement on the fit + evaluate paths.

Design notes

  • Transparent CPU fallback. Every public entry point returns plain np.ndarray / ArrayLike; callers stay device-agnostic. CuPy is imported lazily and gated on availability, so importing the package never requires CUDA.
  • Device selection. On multi-GPU hosts a file-locked counter round-robins across devices to avoid the race where concurrent forks all grab the same "most-free-memory" snapshot. Single-GPU or externally-pinned (CUDA_VISIBLE_DEVICES) hosts short-circuit to device 0.
  • Typing. Optional GPU deps are typed as Any via the [tool.ty] replace-imports-with-any list (cupy, cupyx, numba, kvikio, nvidia), so the type checker stays green without CUDA installed.

Reviewer focus

  • The CPU/GPU agreement in test_gpu_bspline.py is the main automated guard for numerical equivalence — worth a close read since real CUDA isn't exercised in CI.
  • Device selection (__init__.py) uses a fcntl-locked counter file in $TMPDIR; check that the locking/unlocking and the single-GPU fast path match your deployment (e.g. Nextflow setting CUDA_VISIBLE_DEVICES per task).
  • The kvikio/nvcomp paths are CUDA-only and not covered by CI; they're verified manually on the GPU server.

How to verify

uv run ruff check linumpy/gpu/
uv run ty check linumpy/gpu/
uv run pytest linumpy/tests/test_gpu_bspline.py -x -v   # runs the CPU path without CUDA

@FIrgolitsch
FIrgolitsch requested a review from joe-from-mtl as a code owner July 7, 2026 14:58
@FIrgolitsch FIrgolitsch changed the title Rebuild linumpy.gpu core (interface, array_ops, fft_ops, interpolation, morphology, corrections, zarr_io, bspline, bias_field, image_quality) Add GPU acceleration subpackage linumpy.gpu (CuPy + CPU fallback) Jul 8, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant