nibabel icon indicating copy to clipboard operation
nibabel copied to clipboard

Possible feature: add GPU Direct Storage Support for Nibabel

Open yiheng-wang-nv opened this issue 1 year ago • 3 comments

Hi Nibabel maintainers,

This is Yiheng Wang from NVIDIA, I did some research on how to use Kvikio + Nibabel to load .nii images into GPU directly, it helps accelerate the loading time especially on large (100MB+) images. A solution is confirmed work:

  1. using kvikio to load .nii file as bytes from disk into gpu
  2. using nibabel to load headers of the file from disk into cpu
  3. parse bytes according to headers.

Related discussion: https://github.com/nipy/nibabel/discussions/1385

In the Kvikio side, I already added a related tutorial: https://github.com/rapidsai/kvikio/blob/branch-25.04/notebooks/medical_nifti_image_loading_example.ipynb

I'm wondering if it can be a feature directly in nibabel repo. What do you think?

yiheng-wang-nv avatar Dec 11 '24 03:12 yiheng-wang-nv

cc @effigies

yiheng-wang-nv avatar Feb 19 '25 07:02 yiheng-wang-nv

What would your proposed API be?

effigies avatar Feb 19 '25 15:02 effigies

What would your proposed API be?

Hi @effigies , thanks for the question!

My simple thought is to add an arg for nibabel.loadsave.load, like:

load(filename: FileSpec, to_gpu: bool = False, **kwargs)

For gpu loading cases, if calling .get_fdata(), the achieved data will be a cupy array:

img = nib.load('path/to/your/file.nii', to_gpu=True)
data = img.get_fdata()

What do you think?

yiheng-wang-nv avatar Feb 20 '25 03:02 yiheng-wang-nv