Skip to content

CUDA compare and swap with index #8790

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sklam merged 2 commits into from
Mar 8, 2023
Merged

CUDA compare and swap with index #8790

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e4b546b
CUDA compare and swap with index
ianthomas23 Mar 2, 2023
ad4f623
Updates following review
ianthomas23 Mar 6, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions numba/cuda/cudadecl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -514,6 +514,24 @@ def generic(self, args, kws):
return signature(dty, ary, dty, dty)


@register
class Cuda_atomic_cas(AbstractTemplate):
key = cuda.atomic.cas

def generic(self, args, kws):
assert not kws
ary, idx, old, val = args
dty = ary.dtype

if dty not in integer_numba_types:
return

if ary.ndim == 1:
return signature(dty, ary, types.intp, dty, dty)
elif ary.ndim > 1:
return signature(dty, ary, idx, dty, dty)


@register
class Cuda_nanosleep(ConcreteTemplate):
key = cuda.nanosleep
Expand Down Expand Up @@ -602,6 +620,9 @@ def resolve_nanmax(self, mod):
def resolve_compare_and_swap(self, mod):
return types.Function(Cuda_atomic_compare_and_swap)

def resolve_cas(self, mod):
return types.Function(Cuda_atomic_cas)


@register_attr
class CudaFp16Template(AttributeTemplate):
Expand Down
43 changes: 34 additions & 9 deletions numba/cuda/cudaimpl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -708,7 +708,7 @@ def impl(context, builder, sig, args):
lower(math.degrees, types.f8)(gen_deg_rad(_rad2deg))


def _normalize_indices(context, builder, indty, inds):
def _normalize_indices(context, builder, indty, inds, aryty, valty):
"""
Convert integer indices into tuple of intp
"""
Expand All @@ -719,6 +719,15 @@ def _normalize_indices(context, builder, indty, inds):
indices = cgutils.unpack_tuple(builder, inds, count=len(indty))
indices = [context.cast(builder, i, t, types.intp)
for t, i in zip(indty, indices)]

dtype = aryty.dtype
if dtype != valty:
raise TypeError("expect %s but got %s" % (dtype, valty))

if aryty.ndim != len(indty):
raise TypeError("indexing %d-D array with %d-D index" %
(aryty.ndim, len(indty)))

return indty, indices


Expand All @@ -729,14 +738,8 @@ def imp(context, builder, sig, args):
ary, inds, val = args
dtype = aryty.dtype

indty, indices = _normalize_indices(context, builder, indty, inds)

if dtype != valty:
raise TypeError("expect %s but got %s" % (dtype, valty))

if aryty.ndim != len(indty):
raise TypeError("indexing %d-D array with %d-D index" %
(aryty.ndim, len(indty)))
indty, indices = _normalize_indices(context, builder, indty, inds,
aryty, valty)

lary = context.make_array(aryty)(context, builder, ary)
ptr = cgutils.get_item_pointer(context, builder, aryty, lary, indices,
Expand Down Expand Up @@ -935,6 +938,28 @@ def ptx_atomic_cas_tuple(context, builder, sig, args):
'with %s array' % dtype)


@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):

@ianthomas23 ianthomas23 Mar 2, 2023

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@gmarkall gmarkall Mar 6, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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!)

@ianthomas23 ianthomas23 Mar 6, 2023

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, that is much more concise.

@ianthomas23 ianthomas23 Mar 6, 2023

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done.

aryty, indty, oldty, valty = sig.args
ary, inds, old, val = args

indty, indices = _normalize_indices(context, builder, indty, inds, aryty,
valty)

lary = context.make_array(aryty)(context, builder, ary)
ptr = cgutils.get_item_pointer(context, builder, aryty, lary, indices,
wraparound=True)

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 cas with %s array' % aryty.dtype)


# -----------------------------------------------------------------------------

@lower(stubs.nanosleep, types.uint32)
Expand Down
8 changes: 8 additions & 0 deletions numba/cuda/simulator/kernelapi.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -129,6 +129,7 @@ def array(self, shape, dtype):
maxlock = threading.Lock()
minlock = threading.Lock()
caslock = threading.Lock()
casindexlock = threading.Lock()

@gmarkall gmarkall Mar 6, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@ianthomas23 ianthomas23 Mar 6, 2023

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done.

inclock = threading.Lock()
declock = threading.Lock()
exchlock = threading.Lock()
Expand Down Expand Up @@ -221,6 +222,13 @@ def compare_and_swap(self, array, old, val):
array[index] = val
return loaded

def cas(self, array, index, old, val):
with casindexlock:
loaded = array[index]
if loaded == old:
array[index] = val
return loaded


class FakeCUDAFp16(object):
def hadd(self, a, b):
Expand Down
13 changes: 12 additions & 1 deletion numba/cuda/stubs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -503,7 +503,18 @@ class compare_and_swap(Stub):

Supported on int32, int64, uint32, uint64 operands only.

Returns the current value as if it is loaded atomically.
Returns the old value as if it is loaded atomically.
"""

class cas(Stub):
"""cas(ary, idx, old, val)

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

@gmarkall gmarkall Mar 6, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@ianthomas23 ianthomas23 Mar 6, 2023

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done.

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

@gmarkall gmarkall Mar 6, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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``.

@ianthomas23 ianthomas23 Mar 6, 2023

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done.


Supported on int32, int64, uint32, uint64 operands only.

Returns the old value as if it is loaded atomically.
"""


Expand Down
85 changes: 72 additions & 13 deletions numba/cuda/tests/cudapy/test_atomics.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -456,6 +456,20 @@ def atomic_compare_and_swap(res, old, ary, fill_val):
old[gid] = out


def atomic_cas_1dim(res, old, ary, fill_val):
gid = cuda.grid(1)
if gid < res.size:
out = cuda.atomic.cas(res, gid, fill_val, ary[gid])
old[gid] = out

@gmarkall gmarkall Mar 6, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?)

@ianthomas23 ianthomas23 Mar 6, 2023

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done, and on the other two functions too.



def atomic_cas_2dim(res, old, ary, fill_val):
gid = cuda.grid(2)
if gid[0] < res.shape[0] and gid[1] < res.shape[1]:
out = cuda.atomic.cas(res, gid, fill_val, ary[gid])
old[gid] = out

@gmarkall gmarkall Mar 6, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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])



class TestCudaAtomics(CUDATestCase):
def setUp(self):
super().setUp()
Expand Down Expand Up @@ -1251,12 +1265,14 @@ def test_atomic_min_double_shared(self):
gold = np.min(vals)
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):

@ianthomas23 ianthomas23 Mar 2, 2023

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This function is used to test both compare_and_swap and cas.

res = [fill] * (n // 2) + [unfill] * (n // 2)
np.random.shuffle(res)
res = np.asarray(res, dtype=dtype)
if ndim == 2:
res.shape = (10, -1)
out = np.zeros_like(res)
ary = np.random.randint(1, 10, size=res.size).astype(res.dtype)
ary = np.random.randint(1, 10, size=res.shape).astype(res.dtype)

fill_mask = res == fill
unfill_mask = res == unfill
Expand All @@ -1265,33 +1281,76 @@ def check_compare_and_swap(self, n, fill, unfill, dtype):
expect_res[fill_mask] = ary[fill_mask]
expect_res[unfill_mask] = unfill

expect_out = np.zeros_like(out)
expect_out[fill_mask] = res[fill_mask]
expect_out[unfill_mask] = unfill
expect_out = res.copy()

cuda_func = cuda.jit(atomic_compare_and_swap)
cuda_func[10, 10](res, out, ary, fill)
cuda_func = cuda.jit(cas_func)
if ndim == 1:
cuda_func[10, 10](res, out, ary, fill)
else:
cuda_func[(10, 10), (10, 10)](res, out, ary, fill)

np.testing.assert_array_equal(expect_res, res)
np.testing.assert_array_equal(expect_out, out)

def test_atomic_compare_and_swap(self):
self.check_compare_and_swap(n=100, fill=-99, unfill=-1, dtype=np.int32)
self.check_cas(n=100, fill=-99, unfill=-1, dtype=np.int32,
cas_func=atomic_compare_and_swap)

def test_atomic_compare_and_swap2(self):
self.check_compare_and_swap(n=100, fill=-45, unfill=-1, dtype=np.int64)
self.check_cas(n=100, fill=-45, unfill=-1, dtype=np.int64,
cas_func=atomic_compare_and_swap)

def test_atomic_compare_and_swap3(self):
rfill = np.random.randint(50, 500, dtype=np.uint32)
runfill = np.random.randint(1, 25, dtype=np.uint32)
self.check_compare_and_swap(n=100, fill=rfill, unfill=runfill,
dtype=np.uint32)
self.check_cas(n=100, fill=rfill, unfill=runfill, dtype=np.uint32,
cas_func=atomic_compare_and_swap)

def test_atomic_compare_and_swap4(self):
rfill = np.random.randint(50, 500, dtype=np.uint64)
runfill = np.random.randint(1, 25, dtype=np.uint64)
self.check_compare_and_swap(n=100, fill=rfill, unfill=runfill,
dtype=np.uint64)
self.check_cas(n=100, fill=rfill, unfill=runfill, dtype=np.uint64,
cas_func=atomic_compare_and_swap)

def test_atomic_cas_1dim(self):
self.check_cas(n=100, fill=-99, unfill=-1, dtype=np.int32,
cas_func=atomic_cas_1dim)

def test_atomic_cas_2dim(self):
self.check_cas(n=100, fill=-99, unfill=-1, dtype=np.int32,
cas_func=atomic_cas_2dim, ndim=2)

def test_atomic_cas2_1dim(self):
self.check_cas(n=100, fill=-45, unfill=-1, dtype=np.int64,
cas_func=atomic_cas_1dim)

def test_atomic_cas2_2dim(self):
self.check_cas(n=100, fill=-45, unfill=-1, dtype=np.int64,
cas_func=atomic_cas_2dim, ndim=2)

def test_atomic_cas3_1dim(self):
rfill = np.random.randint(50, 500, dtype=np.uint32)
runfill = np.random.randint(1, 25, dtype=np.uint32)
self.check_cas(n=100, fill=rfill, unfill=runfill, dtype=np.uint32,
cas_func=atomic_cas_1dim)

def test_atomic_cas3_2dim(self):
rfill = np.random.randint(50, 500, dtype=np.uint32)
runfill = np.random.randint(1, 25, dtype=np.uint32)
self.check_cas(n=100, fill=rfill, unfill=runfill, dtype=np.uint32,
cas_func=atomic_cas_2dim, ndim=2)

def test_atomic_cas4_1dim(self):
rfill = np.random.randint(50, 500, dtype=np.uint64)
runfill = np.random.randint(1, 25, dtype=np.uint64)
self.check_cas(n=100, fill=rfill, unfill=runfill, dtype=np.uint64,
cas_func=atomic_cas_1dim)

def test_atomic_cas4_2dim(self):
rfill = np.random.randint(50, 500, dtype=np.uint64)
runfill = np.random.randint(1, 25, dtype=np.uint64)
self.check_cas(n=100, fill=rfill, unfill=runfill, dtype=np.uint64,
cas_func=atomic_cas_2dim, ndim=2)

# Tests that the atomic add, min, and max operations return the old value -
# in the simulator, they did not (see Issue #5458). The max and min have
Expand Down