GPU CI

.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

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:

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
 ************************

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",

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.

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