feat: GPU acceleration module and scripts#99
Open
FIrgolitsch wants to merge 5 commits intopr-a-build-toolingfrom
Open
feat: GPU acceleration module and scripts#99FIrgolitsch wants to merge 5 commits intopr-a-build-toolingfrom
FIrgolitsch wants to merge 5 commits intopr-a-build-toolingfrom
Conversation
73f2ced to
ece07fe
Compare
ece07fe to
5aae60d
Compare
e4c127f to
a48129d
Compare
5aae60d to
aa42b00
Compare
a48129d to
82b44be
Compare
aa42b00 to
fdaffcf
Compare
- Delete linumpy/gpu/cuda_env.py entirely - Remove cuda_env imports from gpu/__init__.py - Remove ensure_cuda_env/preload_cuda_libraries from linum_fix_illumination_3d.py and drop the now-obsolete XLA_FLAGS setup (was JAX-specific) - Remove ensure_cuda_env/preload_cuda_libraries from linum_benchmark_gpu.py - Update BaSiCPy docstring to reflect PyTorch backend
ce8e35d to
cc75fc8
Compare
fdaffcf to
5efa698
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PR #99 — GPU Acceleration (unified CPU/GPU scripts)
Adds the
linumpy/gpu/module (CuPy + JAX primitives) and wires the existing CPU scripts in the repo to pick up an optional GPU path via--use_gpu. Stand-alone*_gpu.pyscript variants are intentionally removed in this PR (replaced by the unified CPU scripts in PR #97).New library module
linumpy/gpu/array_ops.py— cupy/jax array primitives and CPU fallbackscorrections.py— GPU illumination / flat-field correctionscuda_env.py— CUDA environment detection (CUDA version, CuPy availability, JAX plugin status) + device selection; used by every GPU-capable script to negotiate the backendfft_ops.py— GPU FFT (phase-correlation building blocks)image_quality.py— GPU variance / sharpness metricsinterpolation.py— GPU affine/linear resamplingmorphology.py— GPU morphological ops (erode/dilate/tophat)normalization.py— GPU intensity normalization (percentile / rolling-ball) paired withlinumpy/preproc/normalization.pyfrom PR feat: utility modules, preprocessing improvements, galvo correction #97registration.py— GPU phase-correlation + refinement; callable fromlinumpy/stitching/registration.pyScripts
scripts/linum_benchmark_gpu.py— benchmark harness comparing CPU vs GPU on the key hot paths (normalization, registration, resampling)scripts/linum_gpu_info.py— report available CUDA / CuPy / JAX devices and plugin statusRemoved (unified into CPU scripts via
--use_gpu)linum_aip_gpu.py,linum_assess_slice_quality_gpu.py,linum_create_mosaic_grid_3d_gpu.py,linum_estimate_transform_gpu.py,linum_fix_illumination_3d_gpu.py,linum_generate_mosaic_aips_gpu.py,linum_normalize_intensities_per_slice_gpu.py,linum_resample_mosaic_grid_gpu.pyShell helpers
shell_scripts/fix_jax_cuda_plugin.sh— resets a broken JAX CUDA plugin install; autodetects CUDA 12 vs 13 and picks the matchingjax-cuda*-pluginwheelCommits
feat: GPU acceleration module and scripts— initiallinumpy/gpu/dropfeat(gpu): unify CPU/GPU scripts, add normalization module— removes the*_gpu.pyscript variants, addsgpu/normalization.py, hardenscuda_env.pyDependencies
gpu/gpu-cuda12/gpu-cuda13optional extras used here--use_gpuand dispatch intolinumpy/gpu/