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?

I think that's used to stop the tests failing here when no GPU is present by not loading any CUDA module ... I'll probably need this, but it might be possible to have the documentation and help text run for the cupy backend, even if the code isn't going to work.

(Clearly the tests do need to pass before merging).

pycbc/fft/backend_cupy.py

pycbc/fft/cupyfft.py

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"

This isn't at the same level. The numpy FFT backend is really bad. It's not clear that's the same for cupy. I've not really seen things limited by the memory allocation. One might want a warning in the scheme initialization that things are not great yet, but I think it doesn't belong here.

pycbc/fft/cupyfft.py

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?

It's probably not ... This is still a prototype.

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

Same as above.

pycbc/types/array_cupy.py

pycbc/vetoes/chisq_cupy.py

pycbc/waveform/utils_cupy.py

spxiwh · 2024-11-25T16:57:33Z

Thanks @GarethCabournDavies I'll respond to some of the things above, but in terms of some of the big picture things:

I don't like the named copyright in PyCBC. This does not accurately reflect contribution. I would prefer if this were removed everywhere with all code copyright of "The PyCBC Team" ... But that's a bigger change.
A number of places you highlight non-existent functions .... There's quite a few more! This is, deliberately, a prototype backend, so things are expected not to exist. Hopefully having it merged will encourage others (ie. you) to fill the gaps.

mj-will · 2024-11-27T10:55:15Z

pycbc/waveform/utils_cupy.py

+    return htilde
+
+
+def fstimeshift(freqseries, phi, kmin, kmax):


kmin and kmax don't appear to be used

I added a FIXME for that ... this function should be converted to an ElementwiseKernel, I think, for performance.

I've added a FIXME here that this block should be changed to a proper ElementwiseKernel using these parameters.

This reverts commit 62a59e5.

spxiwh · 2024-12-02T14:33:29Z

@GarethCabournDavies Is there anything else that should be added at this stage? From my perspective, I would prefer to merge this now. I've added a warning when the scheme is loaded to make it clear that this is still a prototype scheme.

GarethCabournDavies

It looks good for merging to me - one minor question that I'm not 100% on.

The other thing which would be nice to see but not necessary is a note in docs/install_cuda.rst to say that this is available but not mature

GarethCabournDavies · 2024-12-02T14:43:29Z

pycbc/types/frequencyseries.py

+                             delta_f=self.delta_f, epoch=self.epoch,
+                             copy=False)
            tmp[:len(self)] = self[:]

        f = TimeSeries(zeros(tlen,
                           dtype=real_same_precision_as(self)),
-                           delta_t=delta_t)
+                           delta_t=delta_t, copy=False)


Do these changes affect other schemes/backends running?

This is a general optimization improvement.

In the previous version here, we run zeros to generate an array, and then in the FrequencySeries initialization we create another array of zeroes and copy across. There's no reason to copy here as the initial zeros array is not being stored anyway and is otherwise freed. So we should only assign the memory for this new array once, not twice, in all cases.

copy=True is also only partially working on Cupy arrays (in some cases it will fail).

WuShichao · 2024-12-02T16:39:53Z

Does "Y Ddraig Goch" mean there is a dragon? 😄

spxiwh · 2024-12-03T11:14:54Z

I'm merging this now then. I encourage interested folks to propose PRs to improve this backend!

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

spxiwh added 6 commits November 25, 2024 11:20

Gareth's comments

2d7a514

Try to fix test suite

e9e8c66

Refactor LAL constants

89843ea

Probably need BBHx

ea040b2

Move cupy array decleration to when needed

56bf071

Typo

5de37e4

mj-will reviewed Nov 27, 2024

View reviewed changes

spxiwh added 3 commits November 28, 2024 09:52

Don't want time-reversed waveforms!

c21f57a

Add FIXME

7850921

Testing chisq optimization - DO NOT MERGE

62a59e5

GarethCabournDavies mentioned this pull request Nov 29, 2024

Add decompress module for prototype cupy backend #4962

Open

1 task

spxiwh added 2 commits December 2, 2024 08:13

Revert "Testing chisq optimization - DO NOT MERGE"

1c49374

This reverts commit 62a59e5.

Add deprecated warning

34fb40e

GarethCabournDavies approved these changes Dec 2, 2024

View reviewed changes

spxiwh merged commit c06f60a into gwastro:master Dec 3, 2024
29 checks passed

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

spxiwh Nov 25, 2024

spxiwh Nov 25, 2024

GarethCabournDavies Nov 25, 2024

spxiwh Nov 25, 2024

GarethCabournDavies Nov 25, 2024

spxiwh Nov 25, 2024

GarethCabournDavies Nov 25, 2024

spxiwh Nov 25, 2024

spxiwh commented Nov 25, 2024

mj-will Nov 27, 2024

spxiwh Nov 28, 2024

spxiwh Dec 2, 2024

spxiwh commented Dec 2, 2024

GarethCabournDavies left a comment

GarethCabournDavies Dec 2, 2024

spxiwh Dec 2, 2024

WuShichao commented Dec 2, 2024

spxiwh commented Dec 3, 2024

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

Prototype cupy backend #4952

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

spxiwh commented Nov 25, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

spxiwh commented Dec 2, 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

WuShichao commented Dec 2, 2024

spxiwh commented Dec 3, 2024

spxiwh commented Nov 22, 2024 •

edited

Loading