diff --git a/.github/workflows/workflow.yml b/.github/workflows/workflow.yml index 96bfc69195..cf263d0f72 100644 --- a/.github/workflows/workflow.yml +++ b/.github/workflows/workflow.yml @@ -74,3 +74,13 @@ jobs: run: "docs/buildsite.sh" shell: bash + decaf_ampere: + runs-on: self-hosted + steps: + - uses: actions/checkout@v2 + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -e ".[dev,gpu_11x]" + - name: Run + run: pytest diff --git a/README.md b/README.md index 1b46cad8b4..e6cd86f486 100644 --- a/README.md +++ b/README.md @@ -64,13 +64,6 @@ pip install -e ".[dev]" If you prefer not to use Anaconda, or want to manage environments yourself, you should be able to use `pip` with Python >= 3.7. Please see the full documentation for details. -You may optionally install additional packages for GPU extensions: - -``` -# Additional GPU packages (requires CUDA) -pip install -e ".[gpu]" -``` - ### Make sure everything works Once ASPIRE is installed, make sure the unit tests run correctly on your platform by doing: diff --git a/docs/source/installation.rst b/docs/source/installation.rst index f81abe078d..46092615f3 100644 --- a/docs/source/installation.rst +++ b/docs/source/installation.rst @@ -99,6 +99,42 @@ Make sure all unit tests run correctly by doing: Tests currently take around 5 minutes to run, but this depends on your specific machine's resources. +Installing GPU Extensions +************************* + +GPU extensions can be installed using pip. +Extensions are grouped based on CUDA versions. +To find the CUDA driver version, run ``nvidia-smi``. + +.. list-table:: CUDA GPU Extension Versions + :widths: 25 25 + :header-rows: 1 + + * - CUDA Version + - ASPIRE Extension + * - 10.2 + - gpu_102 + * - 11.0 + - gpu_110 + * - 11.1 + - gpu_111 + * - >=11.2 + - gpu_11x + +For example, if you have CUDA 11.7 installed on your system, +the command below would install GPU packages required for ASPIRE. + +:: + + # From PyPI + pip install -e "aspire[gpu_11x]" + + # From a local git repo + pip install -e ".[gpu_11x]" + +By default if GPU extensions are correctly installed, +ASPIRE should automatically begin using the GPU for select components +(such as those using ``nufft``). Generating Documentation ************************ diff --git a/setup.py b/setup.py index a46c7582e5..0714800c1c 100644 --- a/setup.py +++ b/setup.py @@ -47,7 +47,10 @@ def read(fname): # for example gpu packages which may not install for all users, # or developer tools that are handy but not required for users. extras_require={ - "gpu": ["pycuda", "cupy", "cufinufft==1.2"], + "gpu_102": ["pycuda", "cupy-cuda102", "cufinufft==1.2"], + "gpu_110": ["pycuda", "cupy-cuda110", "cufinufft==1.2"], + "gpu_111": ["pycuda", "cupy-cuda111", "cufinufft==1.2"], + "gpu_11x": ["pycuda", "cupy-cuda11x", "cufinufft==1.2"], "dev": [ "black", "bumpversion", diff --git a/src/aspire/config.ini b/src/aspire/config.ini index 251d13dabd..d43a80c952 100644 --- a/src/aspire/config.ini +++ b/src/aspire/config.ini @@ -24,7 +24,7 @@ cg_tol = 1e-5 regularizer = 0. [nfft] -backends = finufft, cufinufft, pynfft +backends = cufinufft, finufft, pynfft [ray] # Ray will default to a OS specific tmp dir. diff --git a/src/aspire/nufft/cufinufft.py b/src/aspire/nufft/cufinufft.py index c2e83dd9a8..dba4cdf4a5 100644 --- a/src/aspire/nufft/cufinufft.py +++ b/src/aspire/nufft/cufinufft.py @@ -7,12 +7,13 @@ from cufinufft import cufinufft from aspire.nufft import Plan +from aspire.utils import complex_type logger = logging.getLogger(__name__) class CufinufftPlan(Plan): - def __init__(self, sz, fourier_pts, epsilon=1e-15, ntransforms=1, **kwargs): + def __init__(self, sz, fourier_pts, epsilon=1e-8, ntransforms=1, **kwargs): """ A plan for non-uniform FFT in 2D or 3D. @@ -27,6 +28,12 @@ def __init__(self, sz, fourier_pts, epsilon=1e-15, ntransforms=1, **kwargs): # Passing "ntransforms" > 1 expects one large higher dimensional array later. self.ntransforms = ntransforms + # Workaround cufinufft A100 singles issue + # ASPIRE-Python/703 + # Cast to doubles. + self._original_dtype = fourier_pts.dtype + fourier_pts = fourier_pts.astype(np.float64, copy=False) + # Basic dtype passthough. dtype = fourier_pts.dtype if dtype == np.float64 or dtype == np.complex128: @@ -95,7 +102,12 @@ def transform(self, signal): `(ntransforms, num_pts)`. """ - if not (signal.dtype == self.dtype or signal.dtype == self.complex_dtype): + # Check we're not forcing a dtype workaround for ASPIRE-Python/703, + # then check if we have a dtype mismatch. + # This avoids false positive complaint for the workaround. + if (self._original_dtype == self.dtype) and not ( + signal.dtype == self.dtype or signal.dtype == self.complex_dtype + ): logger.warning( "Incorrect dtypes passed to (a)nufft." " In the future this will be an error." @@ -131,6 +143,8 @@ def transform(self, signal): self._transform_plan.execute(result_gpu, signal_gpu) result = result_gpu.get() + # ASPIRE-Python/703 + result = result.astype(complex_type(self._original_dtype), copy=False) return result @@ -145,7 +159,12 @@ def adjoint(self, signal): :returns: Transformed signal `(sz)` or `(sz, ntransforms)`. """ - if not (signal.dtype == self.complex_dtype or signal.dtype == self.dtype): + # Check we're not forcing a dtype workaround for ASPIRE-Python/703, + # then check if we have a dtype mismatch. + # This avoids false positive complaint for the workaround. + if (self._original_dtype == self.dtype) and not ( + signal.dtype == self.complex_dtype or signal.dtype == self.dtype + ): logger.warning( "Incorrect dtypes passed to (a)nufft." " In the future this will be an error." @@ -171,5 +190,7 @@ def adjoint(self, signal): self._adjoint_plan.execute(signal_gpu, result_gpu) result = result_gpu.get() + # ASPIRE-Python/703 + result = result.astype(complex_type(self._original_dtype), copy=False) return result