Add JIT examples to CUDA docs by brandon-b-miller · Pull Request #7918 · numba/numba

brandon-b-miller · 2022-03-16T20:03:32Z

Hi! I thought I'd add some examples I have pulled together around some use cases for cuda.jit since there is an issue for it. I added some of my own as well as those I could glean from the discussion in the issue.

There was some discussion of moving these to the numba examples repo but after discussion with @gmarkall we decided to submit them here.

Closes #4266

sklam

Thank you for the contribution.

One worry that I have is ease of testing the validity of the examples as API may drift. I'd suggest include the examples as a runnable .py script and use literalinclude to reference the lines. See example in https://github.com/numba/numba/blame/b57da5b52cb330ec03837a721e49eb8c9b01aa66/docs/source/cuda/examples.rst#L485

gmarkall · 2022-03-29T16:53:51Z

I have looked over the vector example, and in https://github.com/gmarkall/numba/tree/jit-examples in the commit gmarkall@7632af7 I have some suggested edits. The main aim behind the edits is to:

Simplify the wording / shorten it,
Change the example code into a test - the same process should be repeatable for the other examples,
Present more idiomatic usage of CUDA (e.g. not passing in array bounds, using a grid of reasonable size, using device_array_like for empty allocations, etc.)

I'll continue to look at the examples but wanted to push this feedback sooner rather than later.

gmarkall · 2022-03-31T11:14:12Z

+         data[i] = next_temp
+
+         # wait for every thread to write before moving on
+         cuda.syncthreads()


How does syncing at the block level have any meaningful effect here?

It feels like this example needs two copies of the domain along with a cooperative grid sync, or one kernel launch per timestep in order to ensure that all inputs are taken from the previous timestep and not some values coming from the current timestep at the boundaries between blocks.

Am I misunderstanding the operation of the example?

gmarkall · 2022-03-31T11:14:12Z

+         data[i] = next_temp
+
+         # wait for every thread to write before moving on
+         cuda.syncthreads()


How does syncing at the block level have any meaningful effect here?

It feels like this example needs two copies of the domain along with a cooperative grid sync, or one kernel launch per timestep in order to ensure that all inputs are taken from the previous timestep and not some values coming from the current timestep at the boundaries between blocks.

Am I misunderstanding the operation of the example?

Yeah, this example was one that I remembered from academia. I think I derived this from some memory of a pure c++ example which used a grid wide barrier. This will be updated.

gmarkall · 2022-03-31T11:20:07Z

+
+.. _cuda_reduction_shared:
+
+Shared Memory Reduction


We should probably note that this as an example of how to use shared memory rather than the way in which reductions should actually be implemented using Numba - we could point to the @cuda.reduce decorator (https://numba.readthedocs.io/en/stable/cuda/reduction.html) and/or demonstrate its use, as it provides an easier interface for building a reduction and uses a slightly more sophisticated (hopefully more performant) implementation: https://github.com/numba/numba/blob/main/numba/cuda/kernels/reduction.py

I will add a comment here, and I used the cuda.reduce technique later on in the monte-carlo example.

gmarkall · 2022-03-31T11:20:07Z

+
+.. _cuda_reduction_shared:
+
+Shared Memory Reduction


We should probably note that this as an example of how to use shared memory rather than the way in which reductions should actually be implemented using Numba - we could point to the @cuda.reduce decorator (https://numba.readthedocs.io/en/stable/cuda/reduction.html) and/or demonstrate its use, as it provides an easier interface for building a reduction and uses a slightly more sophisticated (hopefully more performant) implementation: https://github.com/numba/numba/blob/main/numba/cuda/kernels/reduction.py

gmarkall · 2022-03-31T11:33:07Z

+
+.. code-block:: python
+
+   print(data_gpu.copy_to_host())


It might be nice to include before and after plots generated with something like:

import matplotlib.pyplot as plt plt.plot(np.arange(len(data)), data_gpu.copy_to_host()) plt.title('Heat in a bar') plt.xlabel('Position') plt.ylabel('Temperature') plt.show()

to give something like:

I really wanted to include some plots here but didn't see any precedent for that elsewhere in the docs. Happy to add them in though.

gmarkall · 2022-03-31T11:33:07Z

+
+.. code-block:: python
+
+   print(data_gpu.copy_to_host())


It might be nice to include before and after plots generated with something like:

import matplotlib.pyplot as plt plt.plot(np.arange(len(data)), data_gpu.copy_to_host()) plt.title('Heat in a bar') plt.xlabel('Position') plt.ylabel('Temperature') plt.show()

to give something like:

gmarkall · 2022-03-31T13:14:14Z

+
+               # check elements 'forward' of this one 
+               # until a new session boundary is found
+               while results[gid + look_ahead] == 0:


How is it guaranteed that if another session boundary is written by the next thread, it is correctly seen by the current thread?

Ahh, you are right. The reason is these used to be two separate kernels that I fused into one, forgetting that the synchronization that comes from two distinct kernel launches is actually critical. The first used to write the boundaries and the second built out the IDs. Will look into alternatives.

gmarkall · 2022-03-31T13:14:14Z

+
+               # check elements 'forward' of this one 
+               # until a new session boundary is found
+               while results[gid + look_ahead] == 0:


How is it guaranteed that if another session boundary is written by the next thread, it is correctly seen by the current thread?

gmarkall · 2022-03-31T13:18:03Z

+      ### DETERMINE SESSION BOUNDARIES
+      is_sess_boundary = 0
+      if 0 < gid < size:
+         if (
+               user_id[gid] != user_id[gid - 1]
+               or timestamp[gid] - timestamp[gid - 1] > session_timeout
+         ):
+               is_sess_boundary = 1
+         else:
+               is_sess_boundary = 0
+      else:
+         is_sess_boundary = 1


Could it be made a little clearer what this is doing by writing it as:

Suggested change

### DETERMINE SESSION BOUNDARIES

is_sess_boundary = 0

if 0 < gid < size:

if (

user_id[gid] != user_id[gid - 1]

or timestamp[gid] - timestamp[gid - 1] > session_timeout

):

is_sess_boundary = 1

else:

is_sess_boundary = 0

else:

is_sess_boundary = 1

# Exit early if our gid is beyond the bounds of the data

if (gid >= len(user_id)):

return

# Determine whether we're on a session boundary

different_user = (user_id[gid] != user_id[gid - 1])

time_diff = timestamp[gid] - timestamp[gid - 1]

timed_out = time_diff > session_timeout

first_datapoint = (gid == 0)

is_sess_boundary = different_user or timed_out or first_datapoint

i.e. the variable names describe the checks and the condition for being a session boundary "reads" in plain language as opposed to being a more complicated expression.

Also this would allow the next if to be removed and the code under it to be moved to the same level as the rest of the function.

(Note that my indentation in the edited code may not match that in the diff at present - there's some weird 3-space indentation going on in places)

What did you think of this suggestion? I think the original version of the code is in the test file.

I had tried this directly and the initial problem was that the 0'th thread can't safely index gid - 1, so it seemed like we can't avoid a special case for that thread. That's how I arrived at the current logic which I think is a mix between the original and this suggestion. I definitely liked the way this logic read more easily so I tried to incorporate that into the approach. Happy to iterate on this though if there's something I am missing.

gmarkall · 2022-03-31T13:18:30Z

+
+   @cuda.jit
+   def sessionize(
+      user_id, timestamp, results, size


Can we not pass in the size and instead use len(user_id) in the kernel?

gmarkall · 2022-03-31T13:18:30Z

+
+   @cuda.jit
+   def sessionize(
+      user_id, timestamp, results, size


Can we not pass in the size and instead use len(user_id) in the kernel?

gmarkall · 2022-03-31T13:32:17Z

+
+.. code-block:: python
+
+   X = cp.array([1, 10, 234])


Although I do love to champion CuPy / Numba interoperability, we should probably keep the Numba examples self-contained and use Numba device arrays here.

gmarkall · 2022-03-31T13:33:07Z

+
+..code-block:: python
+
+   print(business_logic(1,2,3)) # -126.79644737231007


I think the 4-space indentation makes this fail to render as code (but it should probably be converted to a tested example and literalincluded anyway)

gmarkall · 2022-03-31T13:36:16Z

+.. code-block:: python
+
+   @cuda.jit
+   def f(res, xarr, yarr, zarr, size):


Can we use len(xarr) or len(res) instead of passing in the size?

gmarkall

Many thanks for the PR - in addition to the commit I linked to in the previous comment, I've left some additional comments on the diff. I think that each of the examples needs to be converted to a test and literalincluded as @sklam suggests, and this will also highlight a couple of things that prevent the code running as-is.

The use of CuPy is nice but I'd prefer if we kept to Numba code for these examples, although it would be great if we also had a Numba +CuPy interop example (I'm not asking for one to be included in this PR though :-) )

Once the examples are converted to tests I think we might be able to streamline some of the explanations and some of the code too, but I'd prefer to iterate on that once we have the code converted to tested samples.

Co-authored-by: Graham Markall <535640+gmarkall@users.noreply.github.com>

- Use two swappable buffers so that only a single sync is needed - Add a check so that threads outside the domain boundary don't compute - Add a check based on the integral of the solution

gmarkall · 2022-04-28T20:37:23Z

gpuci run tests

The example relies on a single block being launched, but a forall isn't guaranteed to generate a configuration with a single block - for example, on the RTX A6000 it creates the configuration `[2, 768]`. Explicitly configuring the launch resolves this issue.

gmarkall · 2022-04-29T10:31:08Z

gpuci run tests

gmarkall

Many thanks!

gmarkall · 2022-05-11T13:05:48Z

@stuartarchibald / @esc Could this have a CUDA buildfarm run please?

esc · 2022-05-16T13:12:14Z

Build numba_smoketest_cuda_yaml_127 has started

gmarkall · 2022-05-16T21:14:03Z

gpuci run tests

gmarkall · 2022-05-17T08:17:10Z

@esc could this have another buildfarm run please?

esc · 2022-05-17T10:27:37Z

@esc could this have another buildfarm run please?

Build numba_smoketest_cuda_yaml_128 has started

esc · 2022-05-17T13:45:51Z

@esc could this have another buildfarm run please?

Build numba_smoketest_cuda_yaml_128 has started

all green 💚

gmarkall · 2022-05-18T19:37:46Z

all green green_heart

Many thanks!

brandon-b-miller added 2 commits March 11, 2022 09:53

port vecadd example

3c10cff

add kernel examples to cuda docs

b57da5b

brandon-b-miller requested review from esc, sklam and stuartarchibald as code owners March 16, 2022 20:03

stuartarchibald added 3 - Ready for Review Effort - medium Medium size effort needed CUDA CUDA related issue/PR labels Mar 17, 2022

stuartarchibald requested a review from gmarkall March 17, 2022 16:09

sklam reviewed Mar 23, 2022

View reviewed changes

gmarkall reviewed Mar 31, 2022

View reviewed changes

Comment thread docs/source/cuda/examples.rst Outdated

gmarkall requested changes Mar 31, 2022

View reviewed changes

gmarkall added 4 - Waiting on author Waiting for author to respond to review and removed 3 - Ready for Review labels Mar 31, 2022

partially address reviews

f6b30f7

brandon-b-miller and others added 6 commits April 28, 2022 14:44

Apply suggestions from code review

f252db3

Co-authored-by: Graham Markall <535640+gmarkall@users.noreply.github.com>

updates for review

57ed078

Fixes for the Laplace example

daef20d

- Use two swappable buffers so that only a single sync is needed - Add a check so that threads outside the domain boundary don't compute - Add a check based on the integral of the solution

small bugfix for plot=False

abadabc

small fix

8ad3ea0

pacify flake8

4d36f38

gmarkall added 2 commits April 29, 2022 11:29

Skip tests that use CG if cudadevrt is missing

c161e8e

gmarkall previously approved these changes Apr 29, 2022

View reviewed changes

gmarkall added Pending BuildFarm For PRs that have been reviewed but pending a push through our buildfarm 4 - Waiting on CI Review etc done, waiting for CI to finish and removed 4 - Waiting on author Waiting for author to respond to review labels Apr 29, 2022

Skip tests that use grid sync on CC < 6.0

09e1d6d

gmarkall dismissed their stale review via 09e1d6d May 16, 2022 21:13

gmarkall added this to the Numba 0.56 RC milestone May 18, 2022

sklam approved these changes May 23, 2022

View reviewed changes

sklam merged commit ad8fbd9 into numba:main May 23, 2022


		..code-block:: python

		print(business_logic(1,2,3)) # -126.79644737231007

Conversation

brandon-b-miller commented Mar 16, 2022

Uh oh!

sklam left a comment

Choose a reason for hiding this comment

Uh oh!

gmarkall commented Mar 29, 2022

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

gmarkall left a comment

Choose a reason for hiding this comment

Uh oh!

gmarkall commented Apr 28, 2022

Uh oh!

gmarkall commented Apr 29, 2022

Uh oh!

gmarkall left a comment

Choose a reason for hiding this comment

Uh oh!

gmarkall commented May 11, 2022

Uh oh!

esc commented May 16, 2022

Uh oh!

gmarkall commented May 16, 2022

Uh oh!

gmarkall commented May 17, 2022

Uh oh!

esc commented May 17, 2022

Uh oh!

esc commented May 17, 2022

Uh oh!

gmarkall commented May 18, 2022

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!