5.5. Supported Atomic Operations

Numba provides access to some of the atomic operations supported in HSA, in the numba.hsa.atomic class.

5.5.1. Example

The following code demonstrates the use of numba.hsa.atomic.add to count every number in [0,32) occurred in the input array in parallel:

from numba import hsa
import numpy as np

@hsa.jit
def hsa_atomic_histogram(ary):
    tid = hsa.get_local_id(0)
    sm = hsa.shared.array(32, numba.uint32)   # declare shared library
    sm[tid] = 0                               # init values to zero
    hsa.barrier(1)                            # synchronize (wait for init)
    loc = ary[tid] % 32                       # ensure we are in range
    hsa.atomic.add(sm, loc, 1)                # atomic add
    hsa.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