CUDA compare and swap with index by ianthomas23 · Pull Request #8790 · numba/numba

ianthomas23 · 2023-03-02T15:21:33Z

Closes #6702.

Currently numba.cuda.atomic.compare_and_swap only operates on the first array element. This PR adds support for operating on any array element by index. Much of this PR is derived from a previous attempt (#7844) with the permission of the original author @bryevdv. The new function is called numba.cuda.atomic.cas and it takes one more argument (the array index) than compare_and_swap.

My use case is in datashader to create bespoke CUDA atomic operations that are more complicated than simple add or max. Datashader essentially writes to a 2D array, each element of which represents an output pixel, and multiple CUDA threads may write to the same pixel at the same time. The indexed cas function allows the use of a 2D integer array as an array of mutexes, one per pixel, to limit access to a single thread at a time per pixel.

For anyone interested, here is my "datashader lite" implementation in which each pixel is visited multiple times with different integer values. Each pixel stores the two maximum values, in decreasing order, of all visits to that pixel. Without the cas-based locking mechanism the checking and shuffling operation of the max values per pixel isn't atomic.

# Datashader-like max2.
# Locking a single pixel at a time using CUDA CAS mutex.
import cupy
from numba import cuda
import numpy as np

@cuda.jit(device=True)
def lock(mutex, index):  # 2D index here
    while cuda.atomic.cas(mutex, index, 0, 1) != 0:
        pass
    cuda.threadfence()  # This may not be necessary?

@cuda.jit(device=True)
def unlock(mutex, index):  # 2D index here
    cuda.threadfence()
    cuda.atomic.exch(mutex, index, 0)

@cuda.jit
def datashader_lite(x, y, val, max2, mutex):
    offset = cuda.grid(1)
    if offset < x.size:
        i, j, v = x[offset], y[offset], val[offset]
        lock(mutex, (j, i))

        # max aggregator would normally have configurable length, but here hard-coded as 2
        if max2[j, i, 0] == -1:
            max2[j, i, 0] = v
        elif v > max2[j, i, 0]:
            max2[j, i, 1] = max2[j, i, 0]
            max2[j, i, 0] = v
        elif v > max2[j, i, 1]:
            max2[j, i, 1] = v

        unlock(mutex, (j, i))

x   = cupy.array([0, 0, 1, 1, 2, 2,  0,  0,  1, 1, 2, 2])
y   = cupy.array([0, 0, 0, 0, 0, 0,  1,  1,  1, 1, 1, 1])
val = cupy.array([1, 2, 3, 4, 5, 6, 12, 11, 10, 9, 8, 7])

ny, nx = 2, 3  # Pixels in target array
max2 = cupy.full((ny, nx, 2), -1)  # -1 is the unset value
mutex = cupy.zeros((ny, nx), dtype=np.uint32)  # Supports integer dtypes
tpb = 10
bpg = int(np.ceil(len(x)/tpb))
datashader_lite[bpg, tpb](x, y, val, max2, mutex)
print('max2', max2)

Output using this PR is:

max2 [[[ 2  1]
  [ 4  3]
  [ 6  5]]

 [[12 11]
  [10  9]
  [ 8  7]]]

This is my first numba PR 😃

ianthomas23 · 2023-03-02T15:23:47Z

+@lower(stubs.atomic.cas, types.Array, types.intp, types.Any, types.Any)
+@lower(stubs.atomic.cas, types.Array, types.Tuple, types.Any, types.Any)
+@lower(stubs.atomic.cas, types.Array, types.UniTuple, types.Any, types.Any)
+def ptx_atomic_cas(context, builder, sig, args):


This function could probably be combined with the previous one, but I haven't figured out how to do this yet.

One way to combine them would be to modify the original implementation so that it inserts the type of the idx parameter into the signature and a zero value into the args, then calls your new implementation - for example:

diff --git a/numba/cuda/cudaimpl.py b/numba/cuda/cudaimpl.py index 5a81ea7ea..72ed22f53 100644 --- a/numba/cuda/cudaimpl.py +++ b/numba/cuda/cudaimpl.py @@ -920,22 +920,10 @@ def ptx_atomic_nanmin(context, builder, dtype, ptr, val): @lower(stubs.atomic.compare_and_swap, types.Array, types.Any, types.Any) -def ptx_atomic_cas_tuple(context, builder, sig, args): - aryty, oldty, valty = sig.args - ary, old, val = args - dtype = aryty.dtype - - lary = context.make_array(aryty)(context, builder, ary) - zero = context.get_constant(types.intp, 0) - ptr = cgutils.get_item_pointer(context, builder, aryty, lary, (zero,)) - - if aryty.dtype in (cuda.cudadecl.integer_numba_types): - lmod = builder.module - bitwidth = aryty.dtype.bitwidth - return nvvmutils.atomic_cmpxchg(builder, lmod, bitwidth, ptr, old, val) - else: - raise TypeError('Unimplemented atomic compare_and_swap ' - 'with %s array' % dtype) +def ptx_atomic_compare_and_swap(context, builder, sig, args): + sig = sig.return_type(sig.args[0], types.intp, sig.args[1], sig.args[2]) + args = (args[0], context.get_constant(types.intp, 0), args[1], args[2]) + return ptx_atomic_cas(context, builder, sig, args) @lower(stubs.atomic.cas, types.Array, types.intp, types.Any, types.Any)

(note also the change of name of the original function, which apparently made little sense anyway!)

Thanks, that is much more concise.

ianthomas23 · 2023-03-02T15:24:33Z

        np.testing.assert_equal(res, gold)

-    def check_compare_and_swap(self, n, fill, unfill, dtype):
+    def check_cas(self, n, fill, unfill, dtype, cas_func, ndim=1):


This function is used to test both compare_and_swap and cas.

stuartarchibald · 2023-03-06T09:27:33Z

@ianthomas23 Many thanks for the patch. @gmarkall any chance you could review this please?

gmarkall · 2023-03-06T12:28:03Z

gpuci run tests

gmarkall

Many thanks for the PR! This looks good in general, and there are just a few small thoughts on the diff.

The reference docs should probably be updated to mention the new function too, in this section: https://github.com/numba/numba/blob/main/docs/source/cuda-reference/kernel.rst#synchronization-and-atomic-operations

gmarkall · 2023-03-06T12:04:58Z

 maxlock = threading.Lock()
 minlock = threading.Lock()
 caslock = threading.Lock()
+casindexlock = threading.Lock()


For consistency I'd probably call this caslock and rename the existing caslock to compare_and_swaplock.

gmarkall · 2023-03-06T12:06:39Z

+    class cas(Stub):
+        """cas(ary, idx, old, val)
+
+        Conditionally assign ``val`` to the element ary[idx] of an array


Should that refer to the position of the element rather than the value itself?:

Suggested change

Conditionally assign ``val`` to the element ary[idx] of an array

Conditionally assign ``val`` to the element ``idx`` of an array

gmarkall · 2023-03-06T12:07:12Z

+        """cas(ary, idx, old, val)
+
+        Conditionally assign ``val`` to the element ary[idx] of an array
+        ``ary`` if the current value of ary[idx] matches ``old``.


For formatting consistency:

Suggested change

``ary`` if the current value of ary[idx] matches ``old``.

``ary`` if the current value of ``ary[idx]`` matches ``old``.

gmarkall · 2023-03-06T12:15:34Z

+@lower(stubs.atomic.cas, types.Array, types.intp, types.Any, types.Any)
+@lower(stubs.atomic.cas, types.Array, types.Tuple, types.Any, types.Any)
+@lower(stubs.atomic.cas, types.Array, types.UniTuple, types.Any, types.Any)
+def ptx_atomic_cas(context, builder, sig, args):


One way to combine them would be to modify the original implementation so that it inserts the type of the idx parameter into the signature and a zero value into the args, then calls your new implementation - for example:

diff --git a/numba/cuda/cudaimpl.py b/numba/cuda/cudaimpl.py index 5a81ea7ea..72ed22f53 100644 --- a/numba/cuda/cudaimpl.py +++ b/numba/cuda/cudaimpl.py @@ -920,22 +920,10 @@ def ptx_atomic_nanmin(context, builder, dtype, ptr, val): @lower(stubs.atomic.compare_and_swap, types.Array, types.Any, types.Any) -def ptx_atomic_cas_tuple(context, builder, sig, args): - aryty, oldty, valty = sig.args - ary, old, val = args - dtype = aryty.dtype - - lary = context.make_array(aryty)(context, builder, ary) - zero = context.get_constant(types.intp, 0) - ptr = cgutils.get_item_pointer(context, builder, aryty, lary, (zero,)) - - if aryty.dtype in (cuda.cudadecl.integer_numba_types): - lmod = builder.module - bitwidth = aryty.dtype.bitwidth - return nvvmutils.atomic_cmpxchg(builder, lmod, bitwidth, ptr, old, val) - else: - raise TypeError('Unimplemented atomic compare_and_swap ' - 'with %s array' % dtype) +def ptx_atomic_compare_and_swap(context, builder, sig, args): + sig = sig.return_type(sig.args[0], types.intp, sig.args[1], sig.args[2]) + args = (args[0], context.get_constant(types.intp, 0), args[1], args[2]) + return ptx_atomic_cas(context, builder, sig, args) @lower(stubs.atomic.cas, types.Array, types.intp, types.Any, types.Any)

(note also the change of name of the original function, which apparently made little sense anyway!)

gmarkall · 2023-03-06T12:20:08Z

+        out = cuda.atomic.cas(res, gid, fill_val, ary[gid])
+        old[gid] = out


Whilst the original atomic_compare_and_swap uses this pattern, is it a little simpler as just:

Suggested change

out = cuda.atomic.cas(res, gid, fill_val, ary[gid])

old[gid] = out

old[gid] = cuda.atomic.cas(res, gid, fill_val, ary[gid])

(and is it worth updating the original function too?)

Done, and on the other two functions too.

gmarkall · 2023-03-06T12:20:43Z

+        out = cuda.atomic.cas(res, gid, fill_val, ary[gid])
+        old[gid] = out


Similar as above:

Suggested change

out = cuda.atomic.cas(res, gid, fill_val, ary[gid])

old[gid] = out

old[gid] = cuda.atomic.cas(res, gid, fill_val, ary[gid])

ianthomas23 · 2023-03-06T17:48:44Z

The reference docs should probably be updated to mention the new function too, in this section: https://github.com/numba/numba/blob/main/docs/source/cuda-reference/kernel.rst#synchronization-and-atomic-operations

I've added docs for numba.cuda.atomic.cas after the other atomic functions:

I haven't added an entry for compare_and_swap as presumably we should be encouraging use of the new function instead.

ianthomas23 · 2023-03-06T17:51:12Z

    indices for indexing into multiple dimensional arrays. The number of element
    in ``idx`` must match the number of dimension of ``array``.

-    Returns the value of ``array[idx]`` before the storing the new value.


I've removed a few extraneous the from this file, which aren't strictly speaking relevant to this PR. I can revert if preferred.

That's great, many thanks for a nice little tidy-up!

ianthomas23 · 2023-03-06T17:51:54Z

I think I've addressed all review comments so far.

gmarkall · 2023-03-06T18:06:38Z

gpuci run tests

gmarkall

Thanks for the quick updates and extra tidy-up - this is looking great! I think not documenting compare_and_swap was also the right decision - thanks for your thoughtfulness!

I'm going to approve this, with the anticipation that gpuCI passes, at which point it will be ready to merge. I think skipping the buildfarm is fine for this PR since that would only cover Windows in addition to gpuCI, and the implementation overlaps with existing patterns used in the target, I don't anticipate there is much chance of a Windows-specific issue.

CUDA compare and swap with index

e4b546b

ianthomas23 requested a review from gmarkall as a code owner March 2, 2023 15:21

ianthomas23 commented Mar 2, 2023

View reviewed changes

guilhermeleobas added the 3 - Ready for Review label Mar 2, 2023

stuartarchibald added the Effort - medium Medium size effort needed label Mar 6, 2023

stuartarchibald added the CUDA CUDA related issue/PR label Mar 6, 2023

gmarkall requested changes Mar 6, 2023

View reviewed changes

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

Updates following review

ad4f623

ianthomas23 requested review from sklam and stuartarchibald as code owners March 6, 2023 17:41

ianthomas23 commented Mar 6, 2023

View reviewed changes

gmarkall added 4 - Waiting on reviewer Waiting for reviewer to respond to author and removed 4 - Waiting on author Waiting for author to respond to review labels Mar 6, 2023

gmarkall approved these changes Mar 6, 2023

View reviewed changes

gmarkall added 4 - Waiting on CI Review etc done, waiting for CI to finish 5 - Ready to merge Review and testing done, is ready to merge and removed 4 - Waiting on reviewer Waiting for reviewer to respond to author 4 - Waiting on CI Review etc done, waiting for CI to finish labels Mar 6, 2023

gmarkall added this to the Numba 0.57 RC milestone Mar 6, 2023

sklam merged commit 291269b into numba:main Mar 8, 2023

ianthomas23 deleted the 6702_cas branch March 8, 2023 10:12

ianthomas23 mentioned this pull request Mar 22, 2023

Support max_n and min_n reductions on GPU holoviz/datashader#1196

Merged

ianthomas23 mentioned this pull request May 4, 2023

Implement fast CUDA mutex holoviz/datashader#1211

Closed

	Conditionally assign ``val`` to the element ary[idx] of an array
	Conditionally assign ``val`` to the element ``idx`` of an array

	``ary`` if the current value of ary[idx] matches ``old``.
	``ary`` if the current value of ``ary[idx]`` matches ``old``.

		out = cuda.atomic.cas(res, gid, fill_val, ary[gid])
		old[gid] = out

	out = cuda.atomic.cas(res, gid, fill_val, ary[gid])
	old[gid] = out
	old[gid] = cuda.atomic.cas(res, gid, fill_val, ary[gid])

Conversation

ianthomas23 commented Mar 2, 2023

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!

stuartarchibald commented Mar 6, 2023

Uh oh!

gmarkall commented Mar 6, 2023

Uh oh!

gmarkall left a comment

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!

ianthomas23 commented Mar 6, 2023

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ianthomas23 commented Mar 6, 2023

Uh oh!

gmarkall commented Mar 6, 2023

Uh oh!

gmarkall left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants