Add basic multi-gpu support

include/cufinufft_opts.h

@@ -20,6 +20,9 @@ typedef struct cufinufft_opts {   // see cufinufft_default_opts() for defaults
 	int gpu_kerevalmeth; // 0: direct exp(sqrt()), 1: Horner ppval
 
 	int gpu_spreadinterponly; // 0: NUFFT, 1: spread or interpolation only
+
+	/* multi-gpu support */
+	int gpu_device_id;
 } cufinufft_opts;
 
 #endif

python/cufinufft/_cufinufft.py

@@ -87,7 +87,9 @@ def _get_NufftOpts():
         ('gpu_obinsizez', c_int),
         ('gpu_maxsubprobsize', c_int),
         ('gpu_nstreams', c_int),
-        ('gpu_kerevalmeth', c_int)]
+        ('gpu_kerevalmeth', c_int),
+        ('gpu_spreadinterponly', c_int),
+        ('gpu_device_id', c_int)]
     return fields
 
 

python/cufinufft/cufinufft.py

@@ -66,6 +66,9 @@ def __init__(self, nufft_type, modes, n_trans=1, eps=1e-6, isign=None,
             else:
                 isign = +1
 
+        # Need to set the plan here in case something goes wrong later on,
+        # otherwise we error during __del__.
+        self.plan = None
 
         # Setup type bound methods
         self.dtype = np.dtype(dtype)
@@ -98,15 +101,20 @@ def __init__(self, nufft_type, modes, n_trans=1, eps=1e-6, isign=None,
         modes = modes[::-1] + (1,) * (3 - self.dim)
         self.modes = (c_int * 3)(*modes)
 
+        # Get the default option values.
         self.opts = self.default_opts(nufft_type, self.dim)
+
+        # Extract list of valid field names.
+        field_names = [name for name, _ in self.opts._fields_]
+
+        # Assign field names from kwargs if they match up, otherwise error.
         for k, v in kwargs.items():
-            try:
+            if k in field_names:
                 setattr(self.opts, k, v)
-            except AttributeError:
+            else:
                 raise TypeError(f"Invalid option '{k}'")
 
-        # Initialize the plan for this instance
-        self.plan = None
+        # Initialize the plan.
         self._plan()
 
     @staticmethod

python/cufinufft/tests/test_error_checks.py

@@ -52,7 +52,7 @@ def test_set_pts_raises_on_size():
 
     kxyz_gpu = gpuarray.to_gpu(kxyz)
 
-    plan = cufinufft(1, shape, 1, tol, dtype=dtype)
+    plan = cufinufft(1, shape, eps=tol, dtype=dtype)
 
     with pytest.raises(TypeError) as err:
         plan.set_pts(kxyz_gpu[0], kxyz_gpu[1][:4])
@@ -63,6 +63,12 @@ def test_set_pts_raises_on_size():
     assert 'kx and kz must be equal' in err.value.args[0]
 
 
+def test_wrong_field_names():
+    with pytest.raises(TypeError) as err:
+        plan = cufinufft(1, (8, 8), foo="bar")
+    assert "Invalid option 'foo'" in err.value.args[0]
+
+
 def test_exec_raises_on_dtype():
     dtype = np.float32
     complex_dtype = np.complex64

python/cufinufft/tests/test_multi.py

@@ -0,0 +1,69 @@
+import pytest
+
+import numpy as np
+
+import pycuda.driver as drv
+import pycuda.gpuarray as gpuarray
+
+from cufinufft import cufinufft
+
+import utils
+
+
+def test_multi_type1(dtype=np.float32, shape=(16, 16, 16), M=4096, tol=1e-3):
+    complex_dtype = utils._complex_dtype(dtype)
+
+    drv.init()
+
+    dev_count = drv.Device.count()
+
+    if dev_count == 1:
+        pytest.skip()
+
+    devs = [drv.Device(dev_id) for dev_id in range(dev_count)]
+
+    dim = len(shape)
+
+    errs = []
+
+    for dev_id, dev in enumerate(devs):
+        ctx = dev.make_context()
+
+        k = utils.gen_nu_pts(M, dim=dim).astype(dtype)
+        c = utils.gen_nonuniform_data(M).astype(complex_dtype)
+
+        k_gpu = gpuarray.to_gpu(k)
+        c_gpu = gpuarray.to_gpu(c)
+        fk_gpu = gpuarray.GPUArray(shape, dtype=complex_dtype)
+
+        plan = cufinufft(1, shape, eps=tol, dtype=dtype,
+                         gpu_device_id=dev_id)
+
+        plan.set_pts(k_gpu[0], k_gpu[1], k_gpu[2])
+
+        plan.execute(c_gpu, fk_gpu)
+
+        fk = fk_gpu.get()
+
+        ind = int(0.1789 * np.prod(shape))
+
+        fk_est = fk.ravel()[ind]
+        fk_target = utils.direct_type1(c, k, shape, ind)
+
+        type1_rel_err = np.abs(fk_target - fk_est) / np.abs(fk_target)
+
+        print(f'Type 1 relative error (GPU {dev_id}):', type1_rel_err)
+
+        ctx.pop()
+
+        errs.append(type1_rel_err)
+
+    assert all(err < 0.01 for err in errs)
+
+
+def main():
+    test_multi_type1()
+
+
+if __name__ == '__main__':
+    main()

src/cufinufft.cu

@@ -100,7 +100,16 @@ int CUFINUFFT_MAKEPLAN(int type, int dim, int *nmodes, int iflag,
 	doc updated, Barnett, 9/22/20.
 */
 {
-
+        // Mult-GPU support: set the CUDA Device ID:
+        int orig_gpu_device_id;
+        cudaGetDevice(& orig_gpu_device_id);
+        if (opts == NULL) {
+            // options might not be supplied to this function => assume device
+            // 0 by default
+            cudaSetDevice(0);
+        } else {
+            cudaSetDevice(opts->gpu_device_id);
+        }
 
 	cudaEvent_t start, stop;
 	cudaEventCreate(&start);
@@ -267,6 +276,10 @@ int CUFINUFFT_MAKEPLAN(int type, int dim, int *nmodes, int iflag,
 		free(fwkerhalf2);
 	if(dim > 2)
 		free(fwkerhalf3);
+
+        // Multi-GPU support: reset the device ID
+        cudaSetDevice(orig_gpu_device_id);
+
 	return ier;
 }
 
@@ -299,6 +312,12 @@ int CUFINUFFT_SETPTS(int M, FLT* d_kx, FLT* d_ky, FLT* d_kz, int N, FLT *d_s,
 	Melody Shih 07/25/19
 */
 {
+        // Mult-GPU support: set the CUDA Device ID:
+        int orig_gpu_device_id;
+        cudaGetDevice(& orig_gpu_device_id);
+        cudaSetDevice(d_plan->opts.gpu_device_id);
+
+
 	int nf1 = d_plan->nf1;
 	int nf2 = d_plan->nf2;
 	int nf3 = d_plan->nf3;
@@ -363,6 +382,10 @@ int CUFINUFFT_SETPTS(int M, FLT* d_kx, FLT* d_ky, FLT* d_kz, int N, FLT *d_s,
 				if(ier != 0 ){
 					printf("error: cuspread2d_nupts_prop, method(%d)\n",
 						d_plan->opts.gpu_method);
+
+                                        // Multi-GPU support: reset the device ID
+                                        cudaSetDevice(orig_gpu_device_id);
+
 					return 1;
 				}
 			}
@@ -371,6 +394,10 @@ int CUFINUFFT_SETPTS(int M, FLT* d_kx, FLT* d_ky, FLT* d_kz, int N, FLT *d_s,
 				if(ier != 0 ){
 					printf("error: cuspread2d_subprob_prop, method(%d)\n",
 						d_plan->opts.gpu_method);
+
+                                        // Multi-GPU support: reset the device ID
+                                        cudaSetDevice(orig_gpu_device_id);
+
 					return 1;
 				}
 			}
@@ -379,6 +406,10 @@ int CUFINUFFT_SETPTS(int M, FLT* d_kx, FLT* d_ky, FLT* d_kz, int N, FLT *d_s,
 				if(ier != 0 ){
 					printf("error: cuspread2d_paul_prop, method(%d)\n",
 						d_plan->opts.gpu_method);
+
+                                        // Multi-GPU support: reset the device ID
+                                        cudaSetDevice(orig_gpu_device_id);
+
 					return 1;
 				}
 			}
@@ -391,6 +422,10 @@ int CUFINUFFT_SETPTS(int M, FLT* d_kx, FLT* d_ky, FLT* d_kz, int N, FLT *d_s,
 				if(ier != 0 ){
 					printf("error: cuspread3d_blockgather_prop, method(%d)\n",
 						d_plan->opts.gpu_method);
+
+                                        // Multi-GPU support: reset the device ID
+                                        cudaSetDevice(orig_gpu_device_id);
+
 					return ier;
 				}
 			}
@@ -399,6 +434,10 @@ int CUFINUFFT_SETPTS(int M, FLT* d_kx, FLT* d_ky, FLT* d_kz, int N, FLT *d_s,
 				if(ier != 0 ){
 					printf("error: cuspread3d_nuptsdriven_prop, method(%d)\n",
 						d_plan->opts.gpu_method);
+
+                                        // Multi-GPU support: reset the device ID
+                                        cudaSetDevice(orig_gpu_device_id);
+
 					return ier;
 				}
 			}
@@ -407,6 +446,10 @@ int CUFINUFFT_SETPTS(int M, FLT* d_kx, FLT* d_ky, FLT* d_kz, int N, FLT *d_s,
 				if(ier != 0 ){
 					printf("error: cuspread3d_subprob_prop, method(%d)\n",
 						d_plan->opts.gpu_method);
+
+                                        // Multi-GPU support: reset the device ID
+                                        cudaSetDevice(orig_gpu_device_id);
+
 					return ier;
 				}
 			}
@@ -421,6 +464,9 @@ int CUFINUFFT_SETPTS(int M, FLT* d_kx, FLT* d_ky, FLT* d_kz, int N, FLT *d_s,
 		milliseconds/1000);
 #endif
 
+        // Multi-GPU support: reset the device ID
+        cudaSetDevice(orig_gpu_device_id);
+
 	return 0;
 }
 
@@ -444,6 +490,11 @@ int CUFINUFFT_EXECUTE(CUCPX* d_c, CUCPX* d_fk, CUFINUFFT_PLAN d_plan)
 	Melody Shih 07/25/19
 */
 {
+        // Mult-GPU support: set the CUDA Device ID:
+        int orig_gpu_device_id;
+        cudaGetDevice(& orig_gpu_device_id);
+        cudaSetDevice(d_plan->opts.gpu_device_id);
+
 	int ier;
 	int type=d_plan->type;
 	switch(d_plan->dim)
@@ -479,6 +530,10 @@ int CUFINUFFT_EXECUTE(CUCPX* d_c, CUCPX* d_fk, CUFINUFFT_PLAN d_plan)
 		}
 		break;
 	}
+
+        // Multi-GPU support: reset the device ID
+        cudaSetDevice(orig_gpu_device_id);
+
 	return ier;
 }
 
@@ -492,15 +547,23 @@ int CUFINUFFT_DESTROY(CUFINUFFT_PLAN d_plan)
 
 */
 {
+        // Mult-GPU support: set the CUDA Device ID:
+        int orig_gpu_device_id;
+        cudaGetDevice(& orig_gpu_device_id);
+        cudaSetDevice(d_plan->opts.gpu_device_id);
+
 	cudaEvent_t start, stop;
 	cudaEventCreate(&start);
 	cudaEventCreate(&stop);
 
 	cudaEventRecord(start);
 
 	// Can't destroy a Null pointer.
-	if(!d_plan)
+	if(!d_plan) {
+                // Multi-GPU support: reset the device ID
+                cudaSetDevice(orig_gpu_device_id);
 		return 1;
+        }
 
 	if(d_plan->fftplan)
 		cufftDestroy(d_plan->fftplan);
@@ -536,6 +599,8 @@ int CUFINUFFT_DESTROY(CUFINUFFT_PLAN d_plan)
 	/* set pointer to NULL now that we've hopefully free'd the memory. */
 	d_plan = NULL;
 
+        // Multi-GPU support: reset the device ID
+        cudaSetDevice(orig_gpu_device_id);
 	return 0;
 }
 
@@ -613,6 +678,9 @@ int CUFINUFFT_DEFAULT_OPTS(int type, int dim, cufinufft_opts *opts)
 		break;
 	}
 
+        // By default, only use device 0
+        opts->gpu_device_id = 0;
+
 	return 0;
 }
 #ifdef __cplusplus