Supported Atomic Operations¶
Numba provides access to some of the atomic operations supported in HSA, in the
numba.roc.atomic
class.
Example¶
The following code demonstrates the use of numba.roc.atomic.add
to
count every number in [0,32) occurred in the input array in parallel:
from numba import roc
import numpy as np
@roc.jit
def hsa_atomic_histogram(ary):
tid = roc.get_local_id(0)
sm = roc.shared.array(32, numba.uint32) # declare shared library
sm[tid] = 0 # init values to zero
roc.barrier(1) # synchronize (wait for init)
loc = ary[tid] % 32 # ensure we are in range
roc.atomic.add(sm, loc, 1) # atomic add
roc.barrier(1) # synchronize
ary[tid] = sm[tid] # store result inplace
ary = np.random.randint(0, 32, size=32).astype(np.uint32)
orig = ary.copy()
# HSA version
hsa_atomic_histogram[1, 32](ary)
# Expected behavior
gold = np.zeros_like(ary)
for i in range(orig.size):
gold[orig[i]] += 1
print(ary) # HSA kernel result
print(gold) # for comparison