GPU Unit Tests

[garrettwrong]

For the last fews days I have been working on-and-off towards fixing the GPU unit tests. Since they were last maintained a wide variety of failures have developed across the codebase.

The current state of the work is that I have tracked all the unit test failures down to the anufft (adjoint) in single precision.

The following patch will restore all the unit tests at the expense of invoking the anufft (used in evaluate) in doubles. Well, all the unit tests except the 32bit cufinufft adjoint. So that is good.

It seems like the simulated experimental pipeline was working okay (in both single and doubles), so perhaps this appears to only effect some small problems in singles. Not sure yet.

I am focusing down to a small standalone problem, which hopefully we can tackle and fix everything up.

(gpu_utests) [gbwright@caf ASPIRE-Python]$ git diff develop | cat
diff --git a/setup.py b/setup.py
index b1eb913..b5533f8 100644
--- a/setup.py
+++ b/setup.py
@@ -47,7 +47,7 @@ setup(
     #   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"],
         "dev": [
             "black",
             "bumpversion",
diff --git a/src/aspire/nufft/cufinufft.py b/src/aspire/nufft/cufinufft.py
index c2e83dd..56985ab 100644
--- a/src/aspire/nufft/cufinufft.py
+++ b/src/aspire/nufft/cufinufft.py
@@ -72,16 +72,20 @@ class CufinufftPlan(Plan):
             self.ntransforms,
             self.epsilon,
             1,
-            dtype=self.dtype,
+            dtype=np.float64,
             **self.adjoint_opts,
         )
 
         # Note, I store self.fourier_pts_gpu so the GPUArrray life
         #   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.afourier_pts_gpu = self.fourier_pts_gpu
+        if self.dtype == np.float32:
+            self.afourier_pts_gpu = gpuarray.to_gpu(self.fourier_pts.astype(np.float64))
+
+        self._adjoint_plan.set_pts(*self.afourier_pts_gpu)
 
     def transform(self, signal):
         """
@@ -162,11 +166,14 @@ class CufinufftPlan(Plan):
             ), "For multiple transforms, signal stack length should match ntransforms {self.ntransforms}."
             res_shape = (self.ntransforms, *self.sz)
 
         signal_gpu = gpuarray.to_gpu(
-            np.ascontiguousarray(signal, dtype=self.complex_dtype)
+            np.ascontiguousarray(signal, dtype=np.complex128)
         )
 
-        result_gpu = gpuarray.GPUArray(res_shape, dtype=self.complex_dtype)
+        result_gpu = gpuarray.GPUArray(res_shape, dtype=np.complex128)
 
         self._adjoint_plan.execute(signal_gpu, result_gpu)

[garrettwrong]

Updating.

i) I am unable to reproduce as a single test case. These cases run and produce acceptable results.
ii) Commenting out the parameterize_class and manually running each of the cases is successful for both doubles and singles.

I'm thinking maybe one of the single precision cases is leaving in bad state and causing subsequent tests to fail. I will continue by trying to query those sort of details.

[garrettwrong]

I attempted to build cufinufft on caf. The default make which targets older platforms works, but I cannot generate code for these A100 cards (sm_80). Wrote to our friendly admin to see if he can install the correct toolkit (11.x).

In the mean time, will try to build the library on tigergpu and relocate it to PACM (sketchy re linking). I also noticed that PyCUDA is unable to gen code, and I think this is a consequence of the toolkit version mismatch.

I don't know if this is related to the problem I am seeing, but want to i) rule it out ii) have the ability to inject CUDA status checks etc.

Also it should be faster ;)

[garrettwrong]

Just updating this.

A branch with a CUDA C/C++ only example has been pushed to my fork as bug_example several weeks ago.

https://github.com/garrettwrong/cufinufft/compare/master..bug_example

The bug appears as non-deterministic junk for single precision 2d1 transforms on A100 cards. Corruption begins at or before 2d deconv step. I was not able to reproduce on the older generation cards I have at home.

#726 merged in a minimal upcasting work around for ASPIRE.

Marking as wontfix for now because i) my time to work on this was spent ineffectually) ii) FI is working on some major code changes wrt finufft and cufinufft. We can consider picking it back up at a later date.

[garrettwrong]

Noting that this was fixed upstream by the FI team in the lastest FINUFFT work (which includes cuFINUFFT). Upon release and integration of that new version, this workaround should be able to be removed.

[garrettwrong]

Currently testing reverting the A100 workaround after installing cufinufft==2.2.0.

[garrettwrong]

This is complete up to reverting the workaround after another FI release.