Improve GPU-STUMP Speed

[seanlaw]

This follows from a discussion on the Numba discourse channel

Currently, STUMPY can be executed on both the CPUs and the GPUs (i.e., given the same input, both should produce the same results - this is useful for validation). The GPU kernel is essentially self contained in the gpu_stump.py module. In fact, all of the heavy lifting occurs in repeatedly launching a single GPU kernel called _compute_and_update_PI_kernel:

        for i in range(range_start, range_stop):
            _compute_and_update_PI_kernel[blocks_per_grid, threads_per_block](
                i,
                device_T_A,
                device_T_B,
                m,
                device_QT_even,
                device_QT_odd,
                device_QT_first,
                device_M_T,
                device_Σ_T,
                device_μ_Q,
                device_σ_Q,
                k,
                ignore_trivial,
                excl_zone,
                device_profile,
                device_indices,
                True,
            )

STUMPY Installation

Technically, STUMPY can be pip/conda installed but since we are interested in injecting code, let's clone and install from source:

git clone https://github.com/TDAmeritrade/stumpy.git ./stumpy.git
cd stumpy.git
python -m pip install .

For GPUs, it is expected that the appropriate NVIDIA driver and cudatoolkit are also installed.

STUMPY on GPUs

Below is a simple example using a small data set for running the gpu_stump function. gpu_stump is just a light wrapper that does a tiny bit of pre-work and then hands off the computation to _compute_and_update_PI_kernel GPU kernel:

import stumpy
import numpy as np

T = np.random.rand(10000)
m = 50

gpu_mp = stumpy.gpu_stump(T, m)

Validating STUMPY GPU Results

Since we have both CPU and GPU implementations in STUMPY, we can validate that the GPU output is correct using our CPU code (which is appropriate for small to medium sized data that are less than 100K in length):

import stumpy
import numpy as np
import numpy.testing as npt

T = np.random.rand(10000)
m = 50

gpu_mp = stumpy.gpu_stump(T, m)
cpu_mp = stumpy.stump(T, m)

npt.assert_almost_equal(gpu_mp, cpu_mp)

[seanlaw]

@gmarkall Please let know if you need anything else or if there's anything I can add/explain/clarify! Thanks again for your help

[gmarkall]

@seanlaw Thanks for the writeup of steps! I'll take a look at this (but it's very close to the weekend here) so it might be next week before I come back to you.

[seanlaw]

No rush! Please enjoy your time

[gmarkall]

There is nothing preventing asynchronous launches for the kernel - the log shows a long stream of:

...
== CUDA [2056] DEBUG -- call driver api: cuCtxGetCurrent
== CUDA [2056] DEBUG -- call driver api: cuCtxGetDevice
== CUDA [2056] DEBUG -- call driver api: cuCtxGetCurrent
== CUDA [2056] DEBUG -- call driver api: cuCtxGetDevice
== CUDA [2056] DEBUG -- call driver api: cuLaunchKernel
== CUDA [2056] DEBUG -- call driver api: cuCtxGetCurrent
== CUDA [2057] DEBUG -- call driver api: cuCtxGetDevice
== CUDA [2057] DEBUG -- call driver api: cuCtxGetCurrent
== CUDA [2057] DEBUG -- call driver api: cuCtxGetDevice
== CUDA [2057] DEBUG -- call driver api: cuLaunchKernel
...

which is exactly what we'd hope for. One problem is that the kernels are quite short-lived, so the launch overhead dominates the execution time somewhat. I have a Numba PR that reduces the launch overhead (numba/numba#6042), and timing the following (with Xeon Gold 6128, Quadro RTX 8000):

import stumpy
import numpy as np
import numpy.testing as npt
import time

T = np.random.rand(10000)
m = 50

gpu_start = time.perf_counter()
gpu_mp = stumpy.gpu_stump(T, m)
gpu_end = time.perf_counter()

cpu_start = time.perf_counter()
cpu_mp = stumpy.stump(T, m)
cpu_end = time.perf_counter()

npt.assert_almost_equal(gpu_mp, cpu_mp)

print(f"CPU: {cpu_end - cpu_start}")
print(f"GPU: {gpu_end - gpu_start}")

With Numba master:

CPU: 6.643358725006692
GPU: 3.1323348870064365

With PR 6042:

CPU: 6.877583767010947
GPU: 1.2927859439951135

The execution time is reduced by 59% for this size input. When I change the size of T from 10000 to 100000, the GPU times are 31.97s before / 13.44s after. There is more that can be done in Numba to optimise the launch times, but that will (hopefully) be added to future PRs.

Another thought on this is that perhaps if the PR numba/numba#4551 were picked up again so Numba supported cooperative groups, it would be possible to have a global barrier in the kernel that allows the for i in range loop from _gpu_stump to be pushed into the kernel, and then only have one launch in _gpu_stump - this would get rid of a lot of launch overhead, though it would have to be balanced by any limitation that it places on the size of the grid (hopefully not too much for this case).

[seanlaw]

@gmarkall Thanks for coming back to this!

There is nothing preventing asynchronous launches for the kernel - the log shows a long stream of:

Maybe I'm a bit confused/misguided on the terminology here so please correct my understanding if I'm wrong. When you say "asynchronous launches" does that imply "asynchronous kernel execution"? I ask because we do need the for i in range loop to happen in order. So, the i+1th iteration MUST happen only after the ith iteration is completed. If the for loop happens in random order then the results should be wrong and not match the CPU output. The fact that the GPU results match the CPU results seems to imply that the kernels are executed in top down order according to the for i in range loop. Is that correct?

One problem is that the kernels are quite short-lived, so the launch overhead dominates the execution time somewhat. I have a Numba PR that reduces the launch overhead (numba/numba#6042)

So, gpu_stump will typically be used when the data size is in the T = np.random.rand(1_000_000) (1 million) to T = np.random.rand(100_000_000) (100 million+) range. In our very rough performance testing, ~67 million data points takes around 13 days and 6 days to compute on 1x 1080Ti and 2x 1080Ti, respectively. Being able to cut that down in half would be huge for us!

Another thought on this is that perhaps if the PR numba/numba#4551 were picked up again so Numba supported cooperative groups, it would be possible to have a global barrier in the kernel that allows the for i in range loop from _gpu_stump to be pushed into the kernel, and then only have one launch in _gpu_stump - this would get rid of a lot of launch overhead, though it would have to be balanced by any limitation that it places on the size of the grid (hopefully not too much for this case).

I think being able to launch a single kernel and thereby reducing the kernel launch would be super! I have an outstanding issue to write this up for Josh (Patterson) and Ben Zaitlen but I've been looking for a way to demonstrate bigger gains on GPUs compared with CPUs. Perhaps you would be interested in working on this together? I know you're super busy but it would be a really nice story to tell together.

[gmarkall]

Maybe I'm a bit confused/misguided on the terminology here so please correct my understanding if I'm wrong. When you say "asynchronous launches" does that imply "asynchronous kernel execution"? I ask because we do need the for i in range loop to happen in order. So, the i+1th iteration MUST happen only after the ith iteration is completed. If the for loop happens in random order then the results should be wrong and not match the CPU output. The fact that the GPU results match the CPU results seems to imply that the kernels are executed in top down order according to the for i in range loop. Is that correct?

Asynchronous means that they are asynchronous with respect to the CPU, but not each other - they will be queued and executed in order on the default stream - so you are getting the behaviour you need. (To go back to your original Discourse question, streams are for when you don't mind if kernels are asynchronous with respect to each other - you can queue up two kernels in different streams if you don't mind which order they execute in)

(will respond to other two items later, just about to join a meeting)

[seanlaw]

Asynchronous means that they are asynchronous with respect to the CPU, but not each other - they will be queued and executed in order on the default stream - so you are getting the behaviour you need. (To go back to your original Discourse question, streams are for when you don't mind if kernels are asynchronous with respect to each other - you can queue up two kernels in different streams if you don't mind which order they execute in)

Ahhhh! That makes perfect sense. I'm glad I asked the question. 😄 Thank you.

[gmarkall]

So, gpu_stump will typically be used when the data size is in the T = np.random.rand(1_000_000) (1 million) to T = np.random.rand(100_000_000) (100 million+) range. In our very rough performance testing, ~67 million data points takes around 13 days and 6 days to compute on 1x 1080Ti and 2x 1080Ti, respectively. Being able to cut that down in half would be huge for us!

OK - I've done a quick try with N = 1_000_000, and the timings are 325 seconds before, and 135 seconds after - still about 59% faster. Are you able to benchmark on your hardware with PR #6042 to see whether you get similar results?

I think being able to launch a single kernel and thereby reducing the kernel launch would be super! I have an outstanding issue to write this up for Josh (Patterson) and Ben Zaitlen but I've been looking for a way to demonstrate bigger gains on GPUs compared with CPUs. Perhaps you would be interested in working on this together? I know you're super busy but it would be a really nice story to tell together.

I can try putting together something that allows testing this with a cooperative grid group - I'm a little unsure of the exact effort needed for this, and I'm also trying to wrap some things up this week due to being away next week, so I don't know whether it's something I'd get to this week. But I'm generally interested in trying to push on this.

[seanlaw]

OK - I've done a quick try with N = 1_000_000, and the timings are 325 seconds before, and 135 seconds after - still about 59% faster. Are you able to benchmark on your hardware with PR #6042 to see whether you get similar results?

Yes, I would be happy to benchmark it! You'll have to excuse my ignorance here but should I be manually reinstalling Numba by cloning the commit from PR numba/numba#6042? I've only ever installed Numba via pip/conda but never from source. Or do you think I could get away with replacing the three Python modules in my existing installation (i.e., inside of site-packages) with the three files that you had changed?

I can try putting together something that allows testing this with a cooperative grid group - I'm a little unsure of the exact effort needed for this, and I'm also trying to wrap some things up this week due to being away next week, so I don't know whether it's something I'd get to this week. But I'm generally interested in trying to push on this.

Absolutely no rush on this. I'm at your mercy! 😄 I already appreciate everything that you've done so far @gmarkall 🙏

Out of curiosity, aside from Numba itself, is the "cooperative grid group" functionality limited to a particular version of the cudatoolkit? Currently, I'm guessing that most of my users are using something equal to or greater than version 10.1 (basically, people who are using tensorflow and PyTorch). I just want to make sure it's not a bleeding edge capability that would break backwards compatibility for people.

[gmarkall]

A note: I just came across, #64, which also suggests grid sync is required :-)

[gmarkall]

Yes, I would be happy to benchmark it! You'll have to excuse my ignorance here but should I be manually reinstalling Numba by cloning the commit from PR numba/numba#6042? I've only ever installed Numba via pip/conda but never from source. Or do you think I could get away with replacing the three Python modules in my existing installation (i.e., inside of site-packages) with the three files that you had changed?

I would guess the cleanest way to do this is in a new conda environment:

conda create -n numbadev -c numba python=3.8 numpy scipy jinja2 cffi llvmlite
conda activate numbadev
git clone git@github.com:gmarkall/numba.git
cd numba
git checkout grm-native-kernel-launch
python setup.py develop

Then use your normal workflow for building / testing STUMPY. The Numba docs on getting set up for development are here, for reference: https://numba.readthedocs.io/en/stable/developer/contributing.html#getting-set-up

I can try putting together something that allows testing this with a cooperative grid group - I'm a little unsure of the exact effort needed for this, and I'm also trying to wrap some things up this week due to being away next week, so I don't know whether it's something I'd get to this week. But I'm generally interested in trying to push on this.

Absolutely no rush on this. I'm at your mercy! smile I already appreciate everything that you've done so far @gmarkall pray

I've tidied up the CG work and made a PR that is close to being ready: numba/numba#6245

With this PR and a modified version of the kernel containing the loop over i and a grid sync, I get timings:

N         Time (seconds)  Original time (s)
========  ==============  =================
  10000:            0.08               3.13
 100000:            2.21              31.97
1000000:          147.37             325.00

There's some scope for improvement here - altering the block size makes a bit of difference to the timing as well.

Out of curiosity, aside from Numba itself, is the "cooperative grid group" functionality limited to a particular version of the cudatoolkit? Currently, I'm guessing that most of my users are using something equal to or greater than version 10.1 (basically, people who are using tensorflow and PyTorch). I just want to make sure it's not a bleeding edge capability that would break backwards compatibility for people.

CGs are supported since CUDA 9.0, so that shouldn't be a problem. My understanding is that a Pascal or later (CC 6.0 or later) architecture is required for a cooperative grid launch.

[seanlaw]

I would guess the cleanest way to do this is in a new conda environment:
conda create -n numbadev -c numba python=3.8 numpy scipy jinja2 cffi llvmlite
conda activate numbadev
git clone git@github.com:gmarkall/numba.git
cd numba
git checkout grm-native-kernel-launch
python setup.py develop
Then use your normal workflow for building / testing STUMPY. The Numba docs on getting set up for development are here, for reference: https://numba.readthedocs.io/en/stable/developer/contributing.html#getting-set-up

Great! Thank you for the instructions. I will try this out and report back once I've computed the benchmark

With this PR and a modified version of the kernel containing the loop over i and a grid sync, I get timings:

N Time (seconds) Original time (s)
======== ============== =================
10000: 0.08 3.13
100000: 2.21 31.97
1000000: 147.37 325.00
There's some scope for improvement here - altering the block size makes a bit of difference to the timing as well.

Hmmm, your previous timing (with only the changes from numba/numba#6042) showed ~59% improvement and this looks consistent with the N = 1_000_000 result shown in the table directly above (which includes numba/numba#6042, a kernel with the for-loop, and grid sync). Do you think that we are approaching saturation and that the maximum percent increase will plateau around 59%? The results seem to suggest that, for N >= 1_000_000, pushing the for-loop into the kernel and adding grid sync is only helping marginally or, perhaps, not at all. Or maybe I'm interpreting this incorrectly?

CGs are supported since CUDA 9.0, so that shouldn't be a problem. My understanding is that a Pascal or later (CC 6.0 or later) architecture is required for a cooperative grid launch.

That is very useful to know. Thank you!