cufinufft with modeord: 0 CMCL-style, 1 FFT-style

lu1and10 · lu1and10 · commit 5188adac2064 · 2024-04-29T16:12:57.000-04:00
diff --git a/CHANGELOG b/CHANGELOG
@@ -1,6 +1,8 @@
 List of features / changes made / release notes, in reverse chronological order.
 If not stated, FINUFFT is assumed (cuFINUFFT <=1.3 is listed separately).
 
+* cufinufft now supports modeord(type 1,2 only): 0 CMCL-style increasing mode
+  order, 1 FFT-style mode order.
 * CPU plan stage prevents now caps # threads at omp_get_max_threads (being 1
   for single-thread build); warns if this cap was activated (PR 431)
 * new docs troubleshooting accuracy limitations due to condition number of the
diff --git a/include/cufinufft/cudeconvolve.h b/include/cufinufft/cudeconvolve.h
@@ -6,22 +6,22 @@
 namespace cufinufft {
 namespace deconvolve {
 template <typename T>
-__global__ void deconvolve_1d(int ms, int nf1, int fw_width, cuda_complex<T> *fw, cuda_complex<T> *fk, T *fwkerhalf1);
+__global__ void deconvolve_1d(int ms, int nf1, int fw_width, cuda_complex<T> *fw, cuda_complex<T> *fk, T *fwkerhalf1, int modeord);
 template <typename T>
-__global__ void amplify_1d(int ms, int nf1, int fw_width, cuda_complex<T> *fw, cuda_complex<T> *fk, T *fwkerhalf2);
+__global__ void amplify_1d(int ms, int nf1, int fw_width, cuda_complex<T> *fw, cuda_complex<T> *fk, T *fwkerhalf2, int modeord);
 template <typename T>
 __global__ void deconvolve_2d(int ms, int mt, int nf1, int nf2, int fw_width, cuda_complex<T> *fw, cuda_complex<T> *fk,
-                              T *fwkerhalf1, T *fwkerhalf2);
+                              T *fwkerhalf1, T *fwkerhalf2, int modeord);
 template <typename T>
 __global__ void amplify_2d(int ms, int mt, int nf1, int nf2, int fw_width, cuda_complex<T> *fw, cuda_complex<T> *fk,
-                           T *fwkerhalf1, T *fwkerhalf2);
+                           T *fwkerhalf1, T *fwkerhalf2, int modeord);
 
 template <typename T>
 __global__ void deconvolve_3d(int ms, int mt, int mu, int nf1, int nf2, int nf3, int fw_width, cuda_complex<T> *fw,
-                              cuda_complex<T> *fk, T *fwkerhalf1, T *fwkerhalf2, T *fwkerhalf3);
+                              cuda_complex<T> *fk, T *fwkerhalf1, T *fwkerhalf2, T *fwkerhalf3, int modeord);
 template <typename T>
 __global__ void amplify_3d(int ms, int mt, int mu, int nf1, int nf2, int nf3, int fw_width, cuda_complex<T> *fw,
-                           cuda_complex<T> *fk, T *fwkerhalf1, T *fwkerhalf2, T *fwkerhalf3);
+                           cuda_complex<T> *fk, T *fwkerhalf1, T *fwkerhalf2, T *fwkerhalf3, int modeord);
 
 template <typename T>
 int cudeconvolve1d(cufinufft_plan_t<T> *d_mem, int blksize);
diff --git a/include/cufinufft_opts.h b/include/cufinufft_opts.h
@@ -26,6 +26,9 @@ typedef struct cufinufft_opts { // see cufinufft_default_opts() for defaults
     int gpu_device_id;
 
     void *gpu_stream;
+
+    int modeord; // (type 1,2 only): 0 CMCL-style increasing mode order
+                 //                  1 FFT-style mode order
 } cufinufft_opts;
 
 #endif
diff --git a/python/cufinufft/cufinufft/_cufinufft.py b/python/cufinufft/cufinufft/_cufinufft.py
@@ -63,7 +63,8 @@ def _get_NufftOpts():
         ('gpu_spreadinterponly', c_int),
         ('gpu_maxbatchsize', c_int),
         ('gpu_device_id', c_int),
-        ('gpu_stream', c_void_p)
+        ('gpu_stream', c_void_p),
+        ('modeord', c_int)
     ]
     return fields
 
diff --git a/python/cufinufft/cufinufft/_plan.py b/python/cufinufft/cufinufft/_plan.py
@@ -65,7 +65,9 @@ class Plan:
                         memory), ``gpu_sort`` (for ``gpu_method == 1``, 0: no
                         sort, 1: sort), ``gpu_kerevalmeth`` (0: direct
                         exp(sqrt), Horner evaluation), ``gpu_device_id`` (GPU
-                        ID), and ``gpu_stream`` (CUDA stream pointer).
+                        ID), ``gpu_stream`` (CUDA stream pointer) and
+                        ``modeord`` (0: CMCL-compatible mode ordering,
+                        1: FFT-style mode ordering).
     """
 
     def __init__(self, nufft_type, n_modes, n_trans=1, eps=1e-6, isign=None,
diff --git a/python/cufinufft/tests/test_basic.py b/python/cufinufft/tests/test_basic.py
@@ -9,7 +9,7 @@
 # NOTE: Tests below fail for tolerance 1e-4 (error executing plan).
 
 DTYPES = [np.float32, np.float64]
-SHAPES = [(16,), (16, 16), (16, 16, 16)]
+SHAPES = [(16,), (16, 16), (16, 16, 16), (19,), (17, 19), (17, 19, 24)]
 MS = [256, 1024, 4096]
 TOLS = [1e-3, 1e-6]
 OUTPUT_ARGS = [False, True]
@@ -47,6 +47,38 @@ def test_type1(to_gpu, to_cpu, dtype, shape, M, tol, output_arg):
     utils.verify_type1(k, c, fk, tol)
 
 
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("shape", SHAPES)
+@pytest.mark.parametrize("M", MS)
+@pytest.mark.parametrize("tol", TOLS)
+@pytest.mark.parametrize("output_arg", OUTPUT_ARGS)
+def test_type1_modeord(to_gpu, to_cpu, dtype, shape, M, tol, output_arg):
+    complex_dtype = utils._complex_dtype(dtype)
+
+    k, c = utils.type1_problem(dtype, shape, M)
+
+    k_gpu = to_gpu(k)
+    c_gpu = to_gpu(c)
+
+    plan = Plan(1, shape, eps=tol, dtype=complex_dtype, modeord=1)
+
+    # Since k_gpu is an array of shape (dim, M), this will expand to
+    # plan.setpts(k_gpu[0], ..., k_gpu[dim]), allowing us to handle all
+    # dimensions with the same call.
+    plan.setpts(*k_gpu)
+
+    if output_arg:
+        fk_gpu = _compat.array_empty_like(c_gpu, shape, dtype=complex_dtype)
+        plan.execute(c_gpu, out=fk_gpu)
+    else:
+        fk_gpu = plan.execute(c_gpu)
+
+    fk = to_cpu(fk_gpu)
+    fk = np.fft.fftshift(fk)
+
+    utils.verify_type1(k, c, fk, tol)
+
+
 @pytest.mark.parametrize("dtype", DTYPES)
 @pytest.mark.parametrize("shape", SHAPES)
 @pytest.mark.parametrize("M", MS)
diff --git a/src/cuda/cufinufft.cu b/src/cuda/cufinufft.cu
@@ -121,5 +121,7 @@ void cufinufft_default_opts(cufinufft_opts *opts)
 
     // By default, only use device 0
     opts->gpu_device_id = 0;
+
+    opts->modeord = 0;
 }
 }
diff --git a/src/cuda/deconvolve_wrapper.cu b/src/cuda/deconvolve_wrapper.cu
@@ -11,97 +11,97 @@ namespace cufinufft {
 namespace deconvolve {
 /* Kernel for copying fw to fk with amplication by prefac/ker */
 // Note: assume modeord=0: CMCL-compatible mode ordering in fk (from -N/2 up
-// to N/2-1)
+// to N/2-1), modeord=1: FFT-compatible mode ordering in fk (from 0 to N/2-1, then -N/2 up to -1).
 template <typename T>
-__global__ void deconvolve_1d(int ms, int nf1, cuda_complex<T> *fw, cuda_complex<T> *fk, T *fwkerhalf1) {
+__global__ void deconvolve_1d(int ms, int nf1, cuda_complex<T> *fw, cuda_complex<T> *fk, T *fwkerhalf1, int modeord) {
     for (int i = blockDim.x * blockIdx.x + threadIdx.x; i < ms; i += blockDim.x * gridDim.x) {
-        int w1 = i - ms / 2 >= 0 ? i - ms / 2 : nf1 + i - ms / 2;
+        int w1 = ( modeord == 0 ) ? ( (i - ms / 2 >= 0) ? i - ms / 2 : nf1 + i - ms / 2 ) : ( (i - ms + ms / 2 >= 0) ? nf1 + i - ms : i );
 
-        T kervalue = fwkerhalf1[abs(i - ms / 2)];
+        T kervalue = fwkerhalf1[(modeord==0) ? abs(i - ms / 2) : ((i - ms + ms / 2 >= 0) ? ms - i : i)];
         fk[i].x = fw[w1].x / kervalue;
         fk[i].y = fw[w1].y / kervalue;
     }
 }
 
 template <typename T>
 __global__ void deconvolve_2d(int ms, int mt, int nf1, int nf2, cuda_complex<T> *fw, cuda_complex<T> *fk, T *fwkerhalf1,
-                              T *fwkerhalf2) {
+                              T *fwkerhalf2, int modeord) {
     for (int i = blockDim.x * blockIdx.x + threadIdx.x; i < ms * mt; i += blockDim.x * gridDim.x) {
         int k1 = i % ms;
         int k2 = i / ms;
         int outidx = k1 + k2 * ms;
-        int w1 = k1 - ms / 2 >= 0 ? k1 - ms / 2 : nf1 + k1 - ms / 2;
-        int w2 = k2 - mt / 2 >= 0 ? k2 - mt / 2 : nf2 + k2 - mt / 2;
+        int w1 = ( modeord == 0 ) ? ( (k1 - ms / 2 >= 0) ? k1 - ms / 2 : nf1 + k1 - ms / 2 ) : ( (k1 - ms + ms / 2 >= 0) ? nf1 + k1 - ms : k1 );
+        int w2 = ( modeord == 0 ) ? ( (k2 - mt / 2 >= 0) ? k2 - mt / 2 : nf2 + k2 - mt / 2 ) : ( (k2 - mt + mt / 2 >= 0) ? nf2 + k2 - mt : k2 );
         int inidx = w1 + w2 * nf1;
 
-        T kervalue = fwkerhalf1[abs(k1 - ms / 2)] * fwkerhalf2[abs(k2 - mt / 2)];
+        T kervalue = fwkerhalf1[(modeord==0) ? abs(k1 - ms / 2) : ((k1 - ms + ms / 2 >= 0) ? ms - k1 : k1)] * fwkerhalf2[(modeord==0) ? abs(k2 - mt / 2) : ((k2 - mt + mt / 2 >= 0) ? mt - k2 : k2)];
         fk[outidx].x = fw[inidx].x / kervalue;
         fk[outidx].y = fw[inidx].y / kervalue;
     }
 }
 
 template <typename T>
 __global__ void deconvolve_3d(int ms, int mt, int mu, int nf1, int nf2, int nf3, cuda_complex<T> *fw,
-                              cuda_complex<T> *fk, T *fwkerhalf1, T *fwkerhalf2, T *fwkerhalf3) {
+                              cuda_complex<T> *fk, T *fwkerhalf1, T *fwkerhalf2, T *fwkerhalf3, int modeord) {
     for (int i = blockDim.x * blockIdx.x + threadIdx.x; i < ms * mt * mu; i += blockDim.x * gridDim.x) {
         int k1 = i % ms;
         int k2 = (i / ms) % mt;
         int k3 = (i / ms / mt);
         int outidx = k1 + k2 * ms + k3 * ms * mt;
-        int w1 = k1 - ms / 2 >= 0 ? k1 - ms / 2 : nf1 + k1 - ms / 2;
-        int w2 = k2 - mt / 2 >= 0 ? k2 - mt / 2 : nf2 + k2 - mt / 2;
-        int w3 = k3 - mu / 2 >= 0 ? k3 - mu / 2 : nf3 + k3 - mu / 2;
+        int w1 = ( modeord == 0 ) ? ( (k1 - ms / 2 >= 0) ? k1 - ms / 2 : nf1 + k1 - ms / 2 ) : ( (k1 - ms + ms / 2 >= 0) ? nf1 + k1 - ms : k1 );
+        int w2 = ( modeord == 0 ) ? ( (k2 - mt / 2 >= 0) ? k2 - mt / 2 : nf2 + k2 - mt / 2 ) : ( (k2 - mt + mt / 2 >= 0) ? nf2 + k2 - mt : k2 );
+        int w3 = ( modeord == 0 ) ? ( (k3 - mu / 2 >= 0) ? k3 - mu / 2 : nf3 + k3 - mu / 2 ) : ( (k3 - mu + mu / 2 >= 0) ? nf3 + k3 - mu : k3 );
         int inidx = w1 + w2 * nf1 + w3 * nf1 * nf2;
 
-        T kervalue = fwkerhalf1[abs(k1 - ms / 2)] * fwkerhalf2[abs(k2 - mt / 2)] * fwkerhalf3[abs(k3 - mu / 2)];
+        T kervalue = fwkerhalf1[(modeord==0) ? abs(k1 - ms / 2) : ((k1 - ms + ms / 2 >= 0) ? ms - k1 : k1)] * fwkerhalf2[(modeord==0) ? abs(k2 - mt / 2) : ((k2 - mt + mt / 2 >= 0) ? mt - k2 : k2)] * fwkerhalf3[(modeord==0) ? abs(k3 - mu / 2) : ((k3 - mu + mu / 2 >= 0) ? mu - k3 : k3)];
         fk[outidx].x = fw[inidx].x / kervalue;
         fk[outidx].y = fw[inidx].y / kervalue;
     }
 }
 
 /* Kernel for copying fk to fw with same amplication */
 template <typename T>
-__global__ void amplify_1d(int ms, int nf1, cuda_complex<T> *fw, cuda_complex<T> *fk, T *fwkerhalf1) {
+__global__ void amplify_1d(int ms, int nf1, cuda_complex<T> *fw, cuda_complex<T> *fk, T *fwkerhalf1, int modeord) {
     for (int i = blockDim.x * blockIdx.x + threadIdx.x; i < ms; i += blockDim.x * gridDim.x) {
-        int w1 = i - ms / 2 >= 0 ? i - ms / 2 : nf1 + i - ms / 2;
+        int w1 = ( modeord == 0 ) ? ( (i - ms / 2 >= 0) ? i - ms / 2 : nf1 + i - ms / 2 ) : ( (i - ms + ms / 2 >= 0) ? nf1 + i - ms : i );
 
-        T kervalue = fwkerhalf1[abs(i - ms / 2)];
+        T kervalue = fwkerhalf1[(modeord==0) ? abs(i - ms / 2) : ((i - ms + ms / 2 >= 0) ? ms - i : i)];
         fw[w1].x = fk[i].x / kervalue;
         fw[w1].y = fk[i].y / kervalue;
     }
 }
 
 template <typename T>
 __global__ void amplify_2d(int ms, int mt, int nf1, int nf2, cuda_complex<T> *fw, cuda_complex<T> *fk, T *fwkerhalf1,
-                           T *fwkerhalf2) {
+                           T *fwkerhalf2, int modeord) {
     for (int i = blockDim.x * blockIdx.x + threadIdx.x; i < ms * mt; i += blockDim.x * gridDim.x) {
         int k1 = i % ms;
         int k2 = i / ms;
         int inidx = k1 + k2 * ms;
-        int w1 = k1 - ms / 2 >= 0 ? k1 - ms / 2 : nf1 + k1 - ms / 2;
-        int w2 = k2 - mt / 2 >= 0 ? k2 - mt / 2 : nf2 + k2 - mt / 2;
+        int w1 = ( modeord == 0 ) ? ( (k1 - ms / 2 >= 0) ? k1 - ms / 2 : nf1 + k1 - ms / 2 ) : ( (k1 - ms + ms / 2 >= 0) ? nf1 + k1 - ms : k1 );
+        int w2 = ( modeord == 0 ) ? ( (k2 - mt / 2 >= 0) ? k2 - mt / 2 : nf2 + k2 - mt / 2 ) : ( (k2 - mt + mt / 2 >= 0) ? nf2 + k2 - mt : k2 );
         int outidx = w1 + w2 * nf1;
 
-        T kervalue = fwkerhalf1[abs(k1 - ms / 2)] * fwkerhalf2[abs(k2 - mt / 2)];
+        T kervalue = fwkerhalf1[(modeord==0) ? abs(k1 - ms / 2) : ((k1 - ms + ms / 2 >= 0) ? ms - k1 : k1)] * fwkerhalf2[(modeord==0) ? abs(k2 - mt / 2) : ((k2 - mt + mt / 2 >= 0) ? mt - k2 : k2)];
         fw[outidx].x = fk[inidx].x / kervalue;
         fw[outidx].y = fk[inidx].y / kervalue;
     }
 }
 
 template <typename T>
 __global__ void amplify_3d(int ms, int mt, int mu, int nf1, int nf2, int nf3, cuda_complex<T> *fw, cuda_complex<T> *fk,
-                           T *fwkerhalf1, T *fwkerhalf2, T *fwkerhalf3) {
+                           T *fwkerhalf1, T *fwkerhalf2, T *fwkerhalf3, int modeord) {
     for (int i = blockDim.x * blockIdx.x + threadIdx.x; i < ms * mt * mu; i += blockDim.x * gridDim.x) {
         int k1 = i % ms;
         int k2 = (i / ms) % mt;
         int k3 = (i / ms / mt);
         int inidx = k1 + k2 * ms + k3 * ms * mt;
-        int w1 = k1 - ms / 2 >= 0 ? k1 - ms / 2 : nf1 + k1 - ms / 2;
-        int w2 = k2 - mt / 2 >= 0 ? k2 - mt / 2 : nf2 + k2 - mt / 2;
-        int w3 = k3 - mu / 2 >= 0 ? k3 - mu / 2 : nf3 + k3 - mu / 2;
+        int w1 = ( modeord == 0 ) ? ( (k1 - ms / 2 >= 0) ? k1 - ms / 2 : nf1 + k1 - ms / 2 ) : ( (k1 - ms + ms / 2 >= 0) ? nf1 + k1 - ms : k1 );
+        int w2 = ( modeord == 0 ) ? ( (k2 - mt / 2 >= 0) ? k2 - mt / 2 : nf2 + k2 - mt / 2 ) : ( (k2 - mt + mt / 2 >= 0) ? nf2 + k2 - mt : k2 );
+        int w3 = ( modeord == 0 ) ? ( (k3 - mu / 2 >= 0) ? k3 - mu / 2 : nf3 + k3 - mu / 2 ) : ( (k3 - mu + mu / 2 >= 0) ? nf3 + k3 - mu : k3 );
         int outidx = w1 + w2 * nf1 + w3 * nf1 * nf2;
 
-        T kervalue = fwkerhalf1[abs(k1 - ms / 2)] * fwkerhalf2[abs(k2 - mt / 2)] * fwkerhalf3[abs(k3 - mu / 2)];
+        T kervalue = fwkerhalf1[(modeord==0) ? abs(k1 - ms / 2) : ((k1 - ms + ms / 2 >= 0) ? ms - k1 : k1)] * fwkerhalf2[(modeord==0) ? abs(k2 - mt / 2) : ((k2 - mt + mt / 2 >= 0) ? mt - k2 : k2)] * fwkerhalf3[(modeord==0) ? abs(k3 - mu / 2) : ((k3 - mu + mu / 2 >= 0) ? mu - k3 : k3)];
         fw[outidx].x = fk[inidx].x / kervalue;
         fw[outidx].y = fk[inidx].y / kervalue;
     }
@@ -125,13 +125,13 @@ int cudeconvolve1d(cufinufft_plan_t<T> *d_plan, int blksize)
     if (d_plan->spopts.spread_direction == 1) {
         for (int t = 0; t < blksize; t++) {
             deconvolve_1d<<<(nmodes + 256 - 1) / 256, 256, 0, stream>>>(ms, nf1, d_plan->fw + t * nf1,
-                                                                        d_plan->fk + t * nmodes, d_plan->fwkerhalf1);
+                                                                        d_plan->fk + t * nmodes, d_plan->fwkerhalf1, d_plan->opts.modeord);
         }
     } else {
         checkCudaErrors(cudaMemsetAsync(d_plan->fw, 0, maxbatchsize * nf1 * sizeof(cuda_complex<T>), stream));
         for (int t = 0; t < blksize; t++) {
             amplify_1d<<<(nmodes + 256 - 1) / 256, 256, 0, stream>>>(ms, nf1, d_plan->fw + t * nf1,
-                                                                     d_plan->fk + t * nmodes, d_plan->fwkerhalf1);
+                                                                     d_plan->fk + t * nmodes, d_plan->fwkerhalf1, d_plan->opts.modeord);
         }
     }
     return 0;
@@ -158,14 +158,14 @@ int cudeconvolve2d(cufinufft_plan_t<T> *d_plan, int blksize)
         for (int t = 0; t < blksize; t++) {
             deconvolve_2d<<<(nmodes + 256 - 1) / 256, 256, 0, stream>>>(ms, mt, nf1, nf2, d_plan->fw + t * nf1 * nf2,
                                                                         d_plan->fk + t * nmodes, d_plan->fwkerhalf1,
-                                                                        d_plan->fwkerhalf2);
+                                                                        d_plan->fwkerhalf2, d_plan->opts.modeord);
         }
     } else {
         checkCudaErrors(cudaMemsetAsync(d_plan->fw, 0, maxbatchsize * nf1 * nf2 * sizeof(cuda_complex<T>), stream));
         for (int t = 0; t < blksize; t++) {
             amplify_2d<<<(nmodes + 256 - 1) / 256, 256, 0, stream>>>(ms, mt, nf1, nf2, d_plan->fw + t * nf1 * nf2,
                                                                      d_plan->fk + t * nmodes, d_plan->fwkerhalf1,
-                                                                     d_plan->fwkerhalf2);
+                                                                     d_plan->fwkerhalf2, d_plan->opts.modeord);
         }
     }
     return 0;
@@ -193,15 +193,15 @@ int cudeconvolve3d(cufinufft_plan_t<T> *d_plan, int blksize)
         for (int t = 0; t < blksize; t++) {
             deconvolve_3d<<<(nmodes + 256 - 1) / 256, 256, 0, stream>>>(
                 ms, mt, mu, nf1, nf2, nf3, d_plan->fw + t * nf1 * nf2 * nf3, d_plan->fk + t * nmodes,
-                d_plan->fwkerhalf1, d_plan->fwkerhalf2, d_plan->fwkerhalf3);
+                d_plan->fwkerhalf1, d_plan->fwkerhalf2, d_plan->fwkerhalf3, d_plan->opts.modeord);
         }
     } else {
         checkCudaErrors(
             cudaMemsetAsync(d_plan->fw, 0, maxbatchsize * nf1 * nf2 * nf3 * sizeof(cuda_complex<T>), stream));
         for (int t = 0; t < blksize; t++) {
             amplify_3d<<<(nmodes + 256 - 1) / 256, 256, 0, stream>>>(
                 ms, mt, mu, nf1, nf2, nf3, d_plan->fw + t * nf1 * nf2 * nf3, d_plan->fk + t * nmodes,
-                d_plan->fwkerhalf1, d_plan->fwkerhalf2, d_plan->fwkerhalf3);
+                d_plan->fwkerhalf1, d_plan->fwkerhalf2, d_plan->fwkerhalf3, d_plan->opts.modeord);
         }
     }
     return 0;

Original file line number	Diff line number	Diff line change
`@@ -63,7 +63,8 @@ def _get_NufftOpts():`
`63`	`63`	`('gpu_spreadinterponly', c_int),`
`64`	`64`	`('gpu_maxbatchsize', c_int),`
`65`	`65`	`('gpu_device_id', c_int),`
`66`		`- ('gpu_stream', c_void_p)`
	`66`	`+ ('gpu_stream', c_void_p),`
	`67`	`+ ('modeord', c_int)`
`67`	`68`	`]`
`68`	`69`	`return fields`
`69`	`70`
Original file line number	Diff line number	Diff line change
`@@ -121,5 +121,7 @@ void cufinufft_default_opts(cufinufft_opts *opts)`
`121`	`121`
`122`	`122`	`// By default, only use device 0`
`123`	`123`	`opts->gpu_device_id = 0;`
	`124`	`+`
	`125`	`+ opts->modeord = 0;`
`124`	`126`	`}`
`125`	`127`	`}`