Prototype cupy backend #4952

spxiwh · 2024-11-22T13:53:59Z

This adds a prototype CUPY backend to PyCBC.

Our current CUDA GPU backend is not working. There's also a lot more tools now for interacting with CUDA than in 2011. CUPY is really nice, and I think will reduce quite a bit the complexity of our CUDA backend, while still allowing us to use the custom CUDA kernels that exist (as demonstrated in the PR).

This backend will:

Run the premerger likelihood through PyCBC inference (with MPI over multiple cores, but not with openmp).
Mostly run pycbc_inspiral. There's some issue in the chisq module, but I've run out of time to debug it.

I post this now, although I would like to have pycbc_inspiral running before proposing merging ... But I did promise on Wednesday that I would post this.

Others have suggested moving to torch instead. I would like to see a demonstration of this if we want to consider going that route or this one.

ACTIONS

I need to make sure that types are consistent in RawKernel (if not expensive, an explicit check before calling would avoid potential strange errors!)

spxiwh · 2024-11-25T13:48:38Z

>> [Mon 25 Nov 07:45:27 CST 2024] Running pycbc inspiral cupy:openmp with 1 threads


>> [Mon 25 Nov 07:45:48 CST 2024] test for GW150914
Pass: 2 GW150914-like triggers

This is now running the pycbc_inspiral unittest (examples/inspiral) ... It's still probably missing lots of things, and probably isn't well optimized (for inspiral), but I'm happy to get feedback (and potentially merge this) at this point.

GarethCabournDavies

All this looks sensible to me, though I don't feel I can approve yet

The bits that I'd got to look to be the same as what I'd implemented (though I was much slower and hadn't got to certain parts)

Main points I was wanting to ask about:

put your own name down where you have done stuff (even if adding to others')
I've looked where bits have been adapted from and have noticed minor discrepancies that I wasn't sure on, so am asking questions.

GarethCabournDavies · 2024-11-25T14:10:55Z

pycbc/fft/backend_cupy.py

+_backend_dict = {'cupy' : 'cupyfft'}
+_backend_list = ['cupy']
+
+_alist, _adict = _list_available(_backend_list, _backend_dict)


The backend_cuda version of this has the if pycbc.HAVE_CUDA statement and this doesn't. This makes me think, should this backend work when not on a GPU?

GarethCabournDavies · 2024-11-25T14:11:32Z

pycbc/fft/backend_cupy.py

@@ -0,0 +1,37 @@
+#  Copyright (C) 2014 Josh Willis


you wrote this

GarethCabournDavies · 2024-11-25T14:11:43Z

pycbc/fft/cupyfft.py

@@ -0,0 +1,88 @@
+# Copyright (C) 2012  Josh Willis


you wrote this

GarethCabournDavies · 2024-11-25T14:14:36Z

pycbc/fft/cupyfft.py

+    else:
+        raise ValueError(_INV_FFT_MSG.format("IFFT", itype, otype))
+
+


It would be good to have something similar to the numpy warning, i.e "The cupy backend is a prototype, and performance may not be as expected"

GarethCabournDavies · 2024-11-25T14:14:53Z

pycbc/fft/cupyfft.py

+
+class FFT(_BaseFFT):
+    """
+    Class for performing FFTs via the numpy interface.


numpy --> cupy

GarethCabournDavies · 2024-11-25T14:24:55Z

pycbc/types/array_cupy.py

+    if self.dtype == _xp.float32 or self.dtype == _xp.float64:
+        return _xp.argmax(abs(self.data))
+    else:
+        return abs_arg_max_complex(self._data)


I don't see where this is defined?

GarethCabournDavies · 2024-11-25T14:26:03Z

pycbc/types/array_cupy.py

+    if cdtype.kind == 'c':
+        return _xp.sum(self.data.conj() * other, dtype=complex128)
+    else:
+        return inner_real(self.data, other)


same here - i dont see where this is defined

GarethCabournDavies · 2024-11-25T14:27:15Z

pycbc/types/array_cupy.py

+#def numpy(self):
+#    return self._data
+


Suggested change

#def numpy(self):

# return self._data

GarethCabournDavies · 2024-11-25T14:34:06Z

pycbc/vetoes/chisq_cupy.py

+    __shared__ unsigned int s;
+    __shared__ unsigned int e;
+    __shared__ float2 chisq[${NT} * ${NP}];
+    float twopi = 6.283185307179586f;


I don't really like that this is defined here - is there any SSOT we can get this from automatically?

GarethCabournDavies · 2024-11-25T14:50:20Z

pycbc/waveform/utils_cupy.py

+"""This module contains the CPU-specific code for
+   convenience utilities for manipulating waveforms


Suggested change

"""This module contains the CPU-specific code for

convenience utilities for manipulating waveforms

"""This module contains the CuPy-specific code for

convenience utilities for manipulating waveforms

spxiwh added 20 commits November 22, 2024 07:44

Developing GPU CUPY backend

d2a79e6

Don't import aligned!!

f72d346

Adding cupy FFT to scheme

e26afa9

Don't understand why copy here??

d29743a

Get right memory pointer

670ad34

Add CUPY scheme to CLI

a745c8d

Typo in scheme.py

e31dd1b

Remove warning

5e321f3

Adding some more array stuff

dd7bc60

The i is kind of important for a phase shift!!

f0e43f2

Add device num support

cf0d2ce

Set device number

8a76a6d

Add prototype spa_tmplt_cupy

059dc9e

Adding new parts of CUPY backend

1e402ba

Typo in scheme.py

8f5b539

Need to force numpy in PSD estimate

d404116

Add SPATmplt in waveform.py

a273450

Starting to understand RawKernel and CUDA

27090ba

Remove print statement

7c463fa

Force numpy returns

c49ccea

GarethCabournDavies reviewed Nov 25, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Prototype cupy backend #4952

Prototype cupy backend #4952

spxiwh commented Nov 22, 2024 •

edited

Loading

spxiwh commented Nov 25, 2024

GarethCabournDavies left a comment

GarethCabournDavies Nov 25, 2024

GarethCabournDavies Nov 25, 2024

GarethCabournDavies Nov 25, 2024

GarethCabournDavies Nov 25, 2024

GarethCabournDavies Nov 25, 2024

GarethCabournDavies Nov 25, 2024

GarethCabournDavies Nov 25, 2024

GarethCabournDavies Nov 25, 2024

GarethCabournDavies Nov 25, 2024

GarethCabournDavies Nov 25, 2024

		else:
		raise ValueError(_INV_FFT_MSG.format("IFFT", itype, otype))

		"""This module contains the CPU-specific code for
		convenience utilities for manipulating waveforms

Prototype cupy backend #4952

Are you sure you want to change the base?

Prototype cupy backend #4952

Conversation

spxiwh commented Nov 22, 2024 • edited Loading

spxiwh commented Nov 25, 2024

GarethCabournDavies left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

spxiwh commented Nov 22, 2024 •

edited

Loading