Skip to content

Add GPU-accelerated Bluestein FFT for Fresnel propagation - #24

Merged
xiaziyna merged 1 commit into
mainfrom
gpu-bluestein
Apr 3, 2026
Merged

Add GPU-accelerated Bluestein FFT for Fresnel propagation#24
xiaziyna merged 1 commit into
mainfrom
gpu-bluestein

Conversation

@xiaziyna

@xiaziyna xiaziyna commented Apr 3, 2026

Copy link
Copy Markdown
Owner
  • New module: bluestein_fft_gpu.py — PyTorch/cuFFT backend for chunked Bluestein zoom FFT, achieving 15x speedup over CPU (2.5 min vs 60 min on 99k x 99k starshade mask, RTX 3090)
  • CPU v2 optimizations to bluestein_fft.py: dict-based chirp cache, pre-allocated buffer reuse, in-place multiplies, scipy.fft workers=-1
  • GPU auto-detection via use_gpu parameter on FresnelSingle, threaded through simulate_field and propagator.gen_pupil_field
  • Optional dependency: pip install pystarshade[gpu] for torch
  • Fix: N_chunk parameter was ignored (hardcoded to 4) in nchunk_zoom_fresnel_single_fft and simulate_field
  • Tests: 6 new GPU tests + all 38 existing CPU tests pass

- New module: bluestein_fft_gpu.py — PyTorch/cuFFT backend for chunked
  Bluestein zoom FFT, achieving 15x speedup over CPU (2.5 min vs 60 min
  on 99k x 99k starshade mask, RTX 3090)
- CPU v2 optimizations to bluestein_fft.py: dict-based chirp cache,
  pre-allocated buffer reuse, in-place multiplies, scipy.fft workers=-1
- GPU auto-detection via use_gpu parameter on FresnelSingle, threaded
  through simulate_field and propagator.gen_pupil_field
- Optional dependency: pip install pystarshade[gpu] for torch
- Fix: N_chunk parameter was ignored (hardcoded to 4) in
  nchunk_zoom_fresnel_single_fft and simulate_field
- Tests: 6 new GPU tests + all 38 existing CPU tests pass
@xiaziyna
xiaziyna merged commit 7456576 into main Apr 3, 2026
1 check passed
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.

2 participants