diff --git a/.github/workflows/workflow.yml b/.github/workflows/workflow.yml index 4a70116185..2b702a96c9 100644 --- a/.github/workflows/workflow.yml +++ b/.github/workflows/workflow.yml @@ -128,7 +128,7 @@ jobs: - uses: actions/checkout@v3 - name: Install dependencies run: | - pip install -e ".[dev,gpu_12x]" + pip install -e ".[dev,gpu-12x]" - name: Customize config run: | echo "Setup tmp dirs and chmod so others can cleanup." diff --git a/docs/source/installation.rst b/docs/source/installation.rst index 455b4c776e..4a48e3a505 100644 --- a/docs/source/installation.rst +++ b/docs/source/installation.rst @@ -141,24 +141,26 @@ driver version, run ``nvidia-smi`` on the intended system. * - CUDA Version - ASPIRE Extension * - 10.2 - - gpu_102 + - gpu-102 * - 11.0 - - gpu_110 + - gpu-110 * - 11.1 - - gpu_111 + - gpu-111 * - >=11.2 - - gpu_11x + - gpu-11x + * - >=12 + - gpu-12x -For example, if you have CUDA 11.7 installed on your system, +For example, if you have CUDA 12.3 installed on your system, the command below would install GPU packages required for ASPIRE. :: # From a local git repo - pip install -e ".[gpu_11x]" + pip install -e ".[gpu-12x]" # From PyPI - pip install "aspire[gpu_11x]" + pip install "aspire[gpu-12x]" By default if the required GPU extensions are correctly installed, diff --git a/pyproject.toml b/pyproject.toml index 71baee3982..3aa73f5ed6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -63,11 +63,11 @@ dependencies = [ "Source" = "https://github.com/ComputationalCryoEM/ASPIRE-Python" [project.optional-dependencies] -gpu_102 = ["pycuda", "cupy-cuda102", "cufinufft==1.3"] -gpu_110 = ["pycuda", "cupy-cuda110", "cufinufft==1.3"] -gpu_111 = ["pycuda", "cupy-cuda111", "cufinufft==1.3"] -gpu_11x = ["pycuda", "cupy-cuda11x", "cufinufft==1.3"] -gpu_12x = ["pycuda", "cupy-cuda12x", "cufinufft==1.3"] +gpu-102 = ["pycuda", "cupy-cuda102", "cufinufft==1.3"] +gpu-110 = ["pycuda", "cupy-cuda110", "cufinufft==1.3"] +gpu-111 = ["pycuda", "cupy-cuda111", "cufinufft==1.3"] +gpu-11x = ["pycuda", "cupy-cuda11x", "cufinufft==1.3"] +gpu-12x = ["pycuda", "cupy-cuda12x", "cufinufft==2.2.0"] dev = [ "black", "bumpversion", diff --git a/src/aspire/nufft/cufinufft.py b/src/aspire/nufft/cufinufft.py index dba4cdf4a5..465c0b23f9 100644 --- a/src/aspire/nufft/cufinufft.py +++ b/src/aspire/nufft/cufinufft.py @@ -4,7 +4,7 @@ import pycuda.autoinit # noqa: F401 import pycuda.driver as cuda # noqa: F401 import pycuda.gpuarray as gpuarray # noqa: F401 -from cufinufft import cufinufft +from cufinufft import Plan as cufPlan from aspire.nufft import Plan from aspire.utils import complex_type @@ -60,8 +60,8 @@ def __init__(self, sz, fourier_pts, epsilon=1e-8, ntransforms=1, **kwargs): self.num_pts = fourier_pts.shape[1] self.epsilon = max(epsilon, np.finfo(self.dtype).eps) - self._transform_plan = cufinufft( - 2, self.sz, self.ntransforms, self.epsilon, -1, dtype=self.dtype + self._transform_plan = cufPlan( + 2, self.sz, self.ntransforms, self.epsilon, -1, dtype=self.complex_dtype ) self.adjoint_opts = dict() @@ -73,13 +73,13 @@ def __init__(self, sz, fourier_pts, epsilon=1e-8, ntransforms=1, **kwargs): ) self.adjoint_opts["gpu_method"] = 1 - self._adjoint_plan = cufinufft( + self._adjoint_plan = cufPlan( 1, self.sz, self.ntransforms, self.epsilon, 1, - dtype=self.dtype, + dtype=self.complex_dtype, **self.adjoint_opts, ) @@ -87,8 +87,8 @@ def __init__(self, sz, fourier_pts, epsilon=1e-8, ntransforms=1, **kwargs): # is tied to instance, instead of this method. self.fourier_pts_gpu = gpuarray.to_gpu(self.fourier_pts) - self._transform_plan.set_pts(*self.fourier_pts_gpu) - self._adjoint_plan.set_pts(*self.fourier_pts_gpu) + self._transform_plan.setpts(*self.fourier_pts_gpu) + self._adjoint_plan.setpts(*self.fourier_pts_gpu) def transform(self, signal): """ @@ -140,7 +140,7 @@ def transform(self, signal): result_gpu = gpuarray.GPUArray(res_shape, dtype=self.complex_dtype) - self._transform_plan.execute(result_gpu, signal_gpu) + self._transform_plan.execute(signal_gpu, out=result_gpu) result = result_gpu.get() # ASPIRE-Python/703 @@ -187,7 +187,7 @@ def adjoint(self, signal): result_gpu = gpuarray.GPUArray(res_shape, dtype=self.complex_dtype) - self._adjoint_plan.execute(signal_gpu, result_gpu) + self._adjoint_plan.execute(signal_gpu, out=result_gpu) result = result_gpu.get() # ASPIRE-Python/703