CUDA Python: Shared Memory and Atomics with Numba

The first lesson gave every thread its own element, and no thread needed another’s data. Two things change that. Shared memory lets the threads of one block pass values to each other through fast on-chip storage, and atomics let threads from the whole grid update the same location without losing updates. Everything below ran on a Tesla T4 in Google Colab, with CUDA 13.0, Python 3.13, numba 0.61.2 with numba-cuda 0.28.2, and cupy-cuda13x 14.0.1.

A tile in shared memory

Each block loads 256 values into a shared array, then every thread reads the value another thread loaded.

# shared_basics.py - reverse each 256-element tile in shared memory
import numpy as np
from numba import cuda, float32

TILE = 256

@cuda.jit
def reverse_tiles(x, out):
    tile = cuda.shared.array(TILE, float32)  # one copy per block, on the SM
    t = cuda.threadIdx.x
    i = cuda.blockIdx.x * TILE + t
    tile[t] = x[i]                  # every thread loads one element
    cuda.syncthreads()              # wait until the whole tile is loaded
    out[i] = tile[TILE - 1 - t]     # then read an element another thread loaded

n = TILE * 4096
x = np.arange(n, dtype=np.float32)
d_out = cuda.device_array_like(x)
reverse_tiles[n // TILE, TILE](cuda.to_device(x), d_out)
out = d_out.copy_to_host()
expected = x.reshape(-1, TILE)[:, ::-1].ravel()
print("first tile starts with: ", out[:4])
print("second tile starts with:", out[TILE:TILE + 4])
print("matches NumPy:", np.array_equal(out, expected))

$ python shared_basics.py
first tile starts with:  [255. 254. 253. 252.]
second tile starts with: [511. 510. 509. 508.]
matches NumPy: True

cuda.shared.array(TILE, float32) gives each block its own 256-float array, 1 KB of the SM’s on-chip shared memory. Its size must be known when the kernel compiles, so it’s a module constant here. Thread t writes slot t, but reads slot 255 - t, which a thread in a different warp wrote. cuda.syncthreads() makes every thread in the block wait there until all 256 have stored their value. n is an exact multiple of the tile size, which is why the kernel needs no bounds guard.

The same kernel without the barrier

Remove the one cuda.syncthreads() line and launch it three times.

# no_sync.py - the same tile reverse with cuda.syncthreads() removed
import numpy as np
from numba import cuda, float32

TILE = 256

@cuda.jit
def reverse_no_sync(x, out):
    tile = cuda.shared.array(TILE, float32)
    t = cuda.threadIdx.x
    i = cuda.blockIdx.x * TILE + t
    tile[t] = x[i]
    out[i] = tile[TILE - 1 - t]     # no barrier: the other thread may not have written yet

n = TILE * 4096
x = np.arange(n, dtype=np.float32)
expected = x.reshape(-1, TILE)[:, ::-1].ravel()
d_x = cuda.to_device(x)
for launch in range(3):
    d_out = cuda.device_array_like(x)
    reverse_no_sync[n // TILE, TILE](d_x, d_out)
    wrong = np.count_nonzero(d_out.copy_to_host() != expected)
    print(f"launch {launch}: {wrong} of {n} elements wrong")

$ python no_sync.py
launch 0: 429023 of 1048576 elements wrong
launch 1: 433216 of 1048576 elements wrong
launch 2: 431136 of 1048576 elements wrong

About 41 to 42% of the elements came out wrong, and the count changed on every launch: across three runs of the program it was 429,023 to 441,664. A thread reached its read before the thread that owns slot 255 - t had written it, and got whatever that slot still held. That’s a race: the answer depends on timing, so it can even pass a small test by luck. Put the barrier back and the kernel matches NumPy again.

Sum a block in shared memory, then one atomic per block

A tree reduction halves the active threads each step until slot 0 holds the block’s sum. Then one thread per block adds it to the total with cuda.atomic.add.

# block_sum.py - sum 2**24 floats: a tree in shared memory, then one atomic per block
import numpy as np
from numba import cuda, float32

THREADS = 256

@cuda.jit
def block_sum(x, total):
    s = cuda.shared.array(THREADS, float32)
    t = cuda.threadIdx.x
    i = cuda.grid(1)
    s[t] = x[i] if i < x.size else 0.0
    cuda.syncthreads()
    stride = THREADS // 2
    while stride > 0:               # halve the number of active threads each step
        if t < stride:
            s[t] += s[t + stride]
        cuda.syncthreads()
        stride //= 2
    if t == 0:
        cuda.atomic.add(total, 0, s[0])  # one atomic per block, not one per element

rng = np.random.default_rng(0)
x = rng.random(2**24, dtype=np.float32)
d_x = cuda.to_device(x)
blocks = (x.size + THREADS - 1) // THREADS
ref = x.sum(dtype=np.float64)
print("blocks, so atomics:", blocks)
print(f"NumPy float64 sum:  {ref:.2f}")
for launch in range(3):
    d_total = cuda.to_device(np.zeros(1, dtype=np.float32))
    block_sum[blocks, THREADS](d_x, d_total)
    gpu = d_total.copy_to_host()[0]
    print(f"launch {launch}: GPU float32 total {gpu:.2f}, relative error {abs(gpu - ref) / ref:.1e}")

$ python block_sum.py
blocks, so atomics: 65536
NumPy float64 sum:  8388153.94
launch 0: GPU float32 total 8388132.00, relative error 2.6e-06
launch 1: GPU float32 total 8388132.00, relative error 2.6e-06
launch 2: GPU float32 total 8388132.50, relative error 2.6e-06

16.8 million additions became 65,536 atomic additions, one per block; the other work happens in shared memory. Every step needs its own cuda.syncthreads(), because the next step reads slots other threads have just written. Note where it sits: outside the if, so every thread reaches it.

The total is close to NumPy’s float64 sum, 2.4 to 2.6 parts per million off across three runs, but it isn’t identical between launches: 8,388,132.00 twice, then 8,388,132.50, and other runs gave 8,388,133.50. The 65,536 atomic additions happen in whatever order the blocks finish, and float32 rounds differently for each order. Near 8.4 million, float32 values are 0.5 apart, which is why the totals move in steps of 0.5. Practice 4 removes the atomic and gets the same bits every run.

A plain += loses updates

Here about a million threads each add 1 to one counter, once with += and once with cuda.atomic.add.

# atomic_race.py - about a million threads add 1 to the same counter
import numpy as np
from numba import cuda

@cuda.jit
def plain_add(counter):
    counter[0] += 1                 # read, add, write: three separate steps

@cuda.jit
def atomic_add(counter):
    cuda.atomic.add(counter, 0, 1)  # one indivisible read-modify-write

N = 1 << 20
for name, kernel in (("plain +=", plain_add), ("atomic.add", atomic_add)):
    d = cuda.to_device(np.zeros(1, dtype=np.int32))
    kernel[N // 256, 256](d)
    print(f"{name}: counted {d.copy_to_host()[0]} of {N}")

$ python atomic_race.py
plain +=: counted 88 of 1048576
atomic.add: counted 1048576 of 1048576

counter[0] += 1 is three steps: read the value, add 1, write it back. Thousands of threads read the same old value at the same time, and their writes overwrite each other, so only 78 to 96 increments survived across three runs. cuda.atomic.add performs the read, add and write as one step that no other thread can interrupt, so every increment counts. The price is that conflicting atomics on one address queue up, which the next example measures.

A histogram: global atomics vs a shared histogram

Count 16.8 million random bytes into 256 bins: first with one global atomic per element, then with a per-block histogram in shared memory that is merged at the end, launched two ways.

# histogram.py - 256-bin histogram: global atomics vs a per-block shared histogram
import numpy as np
from numba import cuda, int32

BINS = THREADS = 256

@cuda.jit
def hist_global(data, bins):
    i = cuda.grid(1)
    if i < data.size:
        cuda.atomic.add(bins, data[i], 1)        # every element is a global atomic

@cuda.jit
def hist_shared(data, bins):
    local = cuda.shared.array(BINS, int32)
    t = cuda.threadIdx.x
    local[t] = 0                                  # BINS == THREADS: one bin per thread
    cuda.syncthreads()
    for i in range(cuda.grid(1), data.size, cuda.gridsize(1)):  # grid-stride loop
        cuda.atomic.add(local, data[i], 1)        # atomics in on-chip shared memory
    cuda.syncthreads()
    cuda.atomic.add(bins, t, local[t])            # one global atomic per bin per block

def ms_per_launch(kernel, d_data, blocks):
    kernel[blocks, THREADS](d_data, cuda.device_array(BINS, np.int32))  # warm-up compile
    d_bins = cuda.to_device(np.zeros(BINS, dtype=np.int32))
    start, end = cuda.event(), cuda.event()
    start.record()
    for _ in range(10):
        kernel[blocks, THREADS](d_data, d_bins)
    end.record()
    end.synchronize()
    return cuda.event_elapsed_time(start, end) / 10, d_bins.copy_to_host() // 10

rng = np.random.default_rng(1)
data = rng.integers(0, BINS, 2**24, dtype=np.uint8)
d_data = cuda.to_device(data)
ref = np.bincount(data, minlength=BINS)
one_per_thread = (data.size + THREADS - 1) // THREADS
small_grid = 4 * cuda.get_current_device().MULTIPROCESSOR_COUNT    # 4 blocks per SM
runs = (("global atomics", hist_global, one_per_thread, data.size),
        ("shared, one element per thread", hist_shared, one_per_thread, one_per_thread * BINS),
        ("shared, grid-stride", hist_shared, small_grid, small_grid * BINS))
times = []
for name, kernel, blocks, global_atomics in runs:
    ms, bins = ms_per_launch(kernel, d_data, blocks)
    times.append(ms)
    print(f"{name}: {blocks} blocks, {global_atomics:,} global atomics, {ms:.3f} ms, "
          f"matches np.bincount: {np.array_equal(bins, ref)}")
print(f"grid-stride vs global atomics in this run: {times[0] / times[2]:.1f}x faster")

$ python histogram.py
global atomics: 65536 blocks, 16,777,216 global atomics, 9.614 ms, matches np.bincount: True
shared, one element per thread: 65536 blocks, 16,777,216 global atomics, 1.562 ms, matches np.bincount: True
shared, grid-stride: 160 blocks, 40,960 global atomics, 0.171 ms, matches np.bincount: True
grid-stride vs global atomics in this run: 56.1x faster

All three match np.bincount. In hist_global, every element is an atomic on global memory, and all 16.8 million of them fall on only 256 addresses.

The second row runs hist_shared with one element per thread, so it launches 65,536 blocks, and each of those blocks merges its 256 bins with 256 global atomics: 65,536 × 256 is again 16,777,216 global atomics, the same count as hist_global. It was still about 6 to 8 times faster in each of three runs (6.2, 6.2 and 7.8). The difference is where the contention happens. The per-element atomics now go to shared memory, on the SM, where each block competes only with its own 256 threads; the merge atomics are spread over 256 addresses and arrive a block at a time.

The third row launches the same kernel on a small grid: 4 blocks per SM, so 160 blocks on the T4’s 40 SMs. The grid-stride loop, range(cuda.grid(1), data.size, cuda.gridsize(1)), has each thread step through the data in jumps of the whole grid’s width, about 410 elements per thread. Every version merges once per block; what changed is the number of blocks. With only 160 of them, 160 × 256 = 40,960 global atomics remain. That version was about 9 to 10 times faster than the second row and 56 to 75 times faster than hist_global, comparing within each of three runs: 0.127 to 0.171 ms per launch. BINS == THREADS lets each thread clear and merge exactly one bin; Practice 3 handles more bins than threads.

The bins are counted over 10 launches and divided by 10, which works because every launch adds the full histogram again. The first launch, outside the timed loop, compiles the kernel.

Max with atomic.max, and the compare-and-swap loop

Numba has atomics for add, sub, max, min, and/or/xor, exchange and more. For anything else there’s compare-and-swap (CAS), which you can build other atomic operations from.

# atomic_max_cas.py - atomic.max, and the compare-and-swap loop behind custom atomics
import numpy as np
from numba import cuda

@cuda.jit
def max_atomic(values, result):
    i = cuda.grid(1)
    if i < values.size:
        cuda.atomic.max(result, 0, values[i])

@cuda.jit
def max_cas(bits, result_bits):
    i = cuda.grid(1)
    if i < bits.size:
        new = bits[i]
        old = result_bits[0]
        while new > old:                     # only try while we'd raise the max
            seen = cuda.atomic.compare_and_swap(result_bits, old, new)
            if seen == old:                  # nobody changed it in between: we won
                break
            old = seen                       # someone else did: retry against their value

rng = np.random.default_rng(2)
ints = rng.integers(0, 10**9, 2**20, dtype=np.int32)
floats = rng.random(2**20, dtype=np.float32)
for name, arr in (("int32 atomic.max", ints), ("float32 atomic.max", floats)):
    d_res = cuda.to_device(np.zeros(1, dtype=arr.dtype))
    max_atomic[arr.size // 256, 256](cuda.to_device(arr), d_res)
    print(f"{name}: {d_res.copy_to_host()[0]}, np.max: {arr.max()}")
bits = floats.view(np.int32)                 # same bytes, read as int32
d_res = cuda.to_device(np.zeros(1, dtype=np.int32))
max_cas[bits.size // 256, 256](cuda.to_device(bits), d_res)
print(f"float32 via CAS on bits: {d_res.copy_to_host().view(np.float32)[0]}, np.max: {floats.max()}")

$ python atomic_max_cas.py
int32 atomic.max: 999996942, np.max: 999996942
float32 atomic.max: 0.9999992847442627, np.max: 0.9999992847442627
float32 via CAS on bits: 0.9999992847442627, np.max: 0.9999992847442627

All three agree with np.max, and they gave the same result in all three runs, because a maximum doesn’t depend on order the way a float sum does. cuda.atomic.max works directly on the int32 and float32 arrays.

max_cas builds the same operation by hand. compare_and_swap(result_bits, old, new) writes new only if the slot still holds old, and always returns what it found. If that’s old, the swap happened and the thread is done. If not, another thread changed the value in between, so the loop tries again against the newer value, and stops once its own value is no longer bigger.

The kernel compares floats as int32 bits. For non-negative floats, larger values have larger bit patterns, so integer comparison orders them correctly; negative floats don’t follow that rule. Practice 2 counts how often the loop has to retry.

The same histogram as a CuPy RawKernel

CuPy can compile CUDA C++ source directly. Here is the same grid-stride shared-memory histogram, against cupy.bincount on the same data.

# cupy_raw.py - the shared-memory histogram as a CuPy RawKernel, against cupy.bincount
import cupy as cp
import numpy as np

src = r'''
extern "C" __global__ void hist_shared(const unsigned char* data, int n, int* bins) {
    __shared__ int local[256];
    int t = threadIdx.x;
    local[t] = 0;
    __syncthreads();
    for (int i = blockIdx.x * blockDim.x + t; i < n; i += gridDim.x * blockDim.x)
        atomicAdd(&local[data[i]], 1);
    __syncthreads();
    atomicAdd(&bins[t], local[t]);
}
'''
hist = cp.RawKernel(src, "hist_shared")

rng = np.random.default_rng(1)
data = cp.asarray(rng.integers(0, 256, 2**24, dtype=np.uint8))
n = data.size
blocks = 4 * cp.cuda.Device().attributes["MultiProcessorCount"]   # 4 blocks per SM

def raw():
    bins = cp.zeros(256, dtype=cp.int32)
    hist((blocks,), (256,), (data, cp.int32(n), bins))
    return bins

def ms_per_call(fn):
    fn()                                     # warm-up: compiles on first call
    start, end = cp.cuda.Event(), cp.cuda.Event()
    start.record()
    for _ in range(10):
        out = fn()
    end.record()
    end.synchronize()
    return cp.cuda.get_elapsed_time(start, end) / 10, out

ms_raw, b_raw = ms_per_call(raw)
ms_lib, b_lib = ms_per_call(lambda: cp.bincount(data, minlength=256))
print(f"RawKernel shared histogram ({blocks} blocks): {ms_raw:.3f} ms per call")
print(f"cupy.bincount:                          {ms_lib:.3f} ms per call")
print(f"same result: {bool((b_raw == b_lib).all())}, RawKernel {ms_lib / ms_raw:.1f}x faster in this run")

$ python cupy_raw.py
RawKernel shared histogram (160 blocks): 0.210 ms per call
cupy.bincount:                          14.838 ms per call
same result: True, RawKernel 70.7x faster in this run

The same result, and in this case the hand-written kernel won by a wide margin: 0.176 to 0.210 ms against 12.1 to 14.8 ms per call, 69 to 71 times faster within each of three runs. That’s close to the Numba grid-stride version’s 0.127 to 0.171 ms: the same algorithm, compiled from C++ instead of Python. cupy.bincount is a general routine that accepts any number of bins and any non-negative integer input; this lesson didn’t measure where its time goes. The kernel knows there are exactly 256 bins and that each fits in one block’s shared memory. That’s the usual trade: write a custom kernel when you know something about your data that a general library call can’t assume, and measure it, because the library is the better default otherwise.

Inside the RawKernel, __shared__ int local[256] is the C++ form of cuda.shared.array, __syncthreads() of cuda.syncthreads(), and atomicAdd of cuda.atomic.add. The C++ lessons cover them in that syntax.

Common gotchas

• A cuda.syncthreads() inside a branch some threads skip can hang or corrupt the block: every thread in the block must reach the same barrier. Keep barriers outside if statements that depend on the thread index, as the block sum does.
• The shared array’s size must be a compile-time constant. Use a module-level constant, or cuda.shared.array(0, dtype) with the size given at launch (Practice 1).
• Shared memory is limited: by default 48 KB per block in Numba-compiled kernels on this T4, and the compiler refuses more (Practice 3).
• A shared array starts with whatever was there before. Clear it before you accumulate into it, as hist_shared does, and put a barrier after the clearing.
• Atomics are correct, but they cost time. Many threads hitting one address queue up behind each other. Reduce in shared memory first, then make one atomic per block.

Bank conflicts, padding and warp shuffles, which Numba also exposes as cuda.shfl_*, are taught in CUDA C++: shared memory and warps. The profiler comparison of reduction strategies is in CUDA C++: atomics and reductions. Why a shared-memory tile doesn’t always beat the cache is in the CUDA memory hierarchy.

Practice

1. Make the tile size a launch-time choice: declare the shared array with size 0, and pass its size in bytes as the fourth launch value. Run the tile reverse with 128 and 512 threads per block.

# p1_dynamic.py - the tile reverse with the tile size chosen at launch time
import numpy as np
from numba import cuda, float32

@cuda.jit
def reverse_tiles(x, out):
    tile = cuda.shared.array(0, float32)     # size 0: the launch decides it
    t = cuda.threadIdx.x
    size = cuda.blockDim.x
    i = cuda.blockIdx.x * size + t
    tile[t] = x[i]
    cuda.syncthreads()
    out[i] = tile[size - 1 - t]

x = np.arange(1 << 20, dtype=np.float32)
for size in (128, 512):
    d_out = cuda.device_array_like(x)
    # 4th launch value: bytes of dynamic shared memory per block
    reverse_tiles[x.size // size, size, 0, size * 4](cuda.to_device(x), d_out)
    expected = x.reshape(-1, size)[:, ::-1].ravel()
    print(f"tile {size}: matches NumPy: {np.array_equal(d_out.copy_to_host(), expected)}")

$ python p1_dynamic.py
tile 128: matches NumPy: True
tile 512: matches NumPy: True

2. Count how often the compare-and-swap loop retries, with random input and with the input sorted in ascending order.

# p2_retries.py - count how often the compare-and-swap loop has to retry
import numpy as np
from numba import cuda

@cuda.jit
def max_cas(bits, result_bits, retries):
    i = cuda.grid(1)
    if i < bits.size:
        new = bits[i]
        old = result_bits[0]
        while new > old:
            seen = cuda.atomic.compare_and_swap(result_bits, old, new)
            if seen == old:
                break
            cuda.atomic.add(retries, 0, 1)   # someone else changed it first
            old = seen

rng = np.random.default_rng(2)
for name, floats in (("random order", rng.random(2**20, dtype=np.float32)),
                     ("ascending", np.sort(rng.random(2**20, dtype=np.float32)))):
    d_res = cuda.to_device(np.zeros(1, dtype=np.int32))
    d_retries = cuda.to_device(np.zeros(1, dtype=np.int32))
    max_cas[floats.size // 256, 256](cuda.to_device(floats.view(np.int32)), d_res, d_retries)
    ok = d_res.copy_to_host().view(np.float32)[0] == floats.max()
    print(f"{name}: correct max: {ok}, retries: {d_retries.copy_to_host()[0]}")

$ python p2_retries.py
random order: correct max: True, retries: 75894
ascending: correct max: True, retries: 855073

3. Make the shared histogram work with more bins than threads, using a loop so each thread clears and merges several bins. Try 1024 bins, then 16,384.

# p3_bins.py - a shared histogram with more bins than threads, until it stops fitting
import re
import numpy as np
from numba import cuda, int32

THREADS = 256

def make_kernel(bins):
    @cuda.jit
    def hist(data, out):
        local = cuda.shared.array(bins, int32)
        t = cuda.threadIdx.x
        for b in range(t, bins, THREADS):     # each thread clears several bins
            local[b] = 0
        cuda.syncthreads()
        i = cuda.grid(1)
        if i < data.size:
            cuda.atomic.add(local, data[i], 1)
        cuda.syncthreads()
        for b in range(t, bins, THREADS):
            cuda.atomic.add(out, b, local[b])
    return hist

rng = np.random.default_rng(1)
for bins in (1024, 16384):
    data = rng.integers(0, bins, 2**22).astype(np.int32)
    d_out = cuda.to_device(np.zeros(bins, dtype=np.int32))
    try:
        make_kernel(bins)[(data.size + THREADS - 1) // THREADS, THREADS](cuda.to_device(data), d_out)
        ok = np.array_equal(d_out.copy_to_host(), np.bincount(data, minlength=bins))
        print(f"{bins} bins ({bins * 4 // 1024} KB shared): matches np.bincount: {ok}")
    except Exception as e:                    # ptxas refuses it when the kernel compiles
        reason = re.search(r"uses too much shared data \([^)]*\)", str(e)).group(0)
        print(f"{bins} bins ({bins * 4 // 1024} KB shared): {type(e).__name__}: {reason}")

$ python p3_bins.py
1024 bins (4 KB shared): matches np.bincount: True
16384 bins (64 KB shared): nvJitLinkError: uses too much shared data (0x10000 bytes, 0xc000 max)

4. Replace the block sum’s atomic with a second kernel pass: each block writes its partial sum to its own slot, and the same kernel then reduces the partials until one value is left. Run it three times.

# p4_two_pass.py - the block sum without atomics: write partials, then reduce them again
import numpy as np
from numba import cuda, float32

THREADS = 256

@cuda.jit
def partial_sums(x, partials):
    s = cuda.shared.array(THREADS, float32)
    t = cuda.threadIdx.x
    i = cuda.grid(1)
    s[t] = x[i] if i < x.size else 0.0
    cuda.syncthreads()
    stride = THREADS // 2
    while stride > 0:
        if t < stride:
            s[t] += s[t + stride]
        cuda.syncthreads()
        stride //= 2
    if t == 0:
        partials[cuda.blockIdx.x] = s[0]   # a plain write: one slot per block

def total(d_x):
    while d_x.size > 1:                    # 2**24 -> 65536 -> 256 -> 1 values
        blocks = (d_x.size + THREADS - 1) // THREADS
        d_next = cuda.device_array(blocks, np.float32)
        partial_sums[blocks, THREADS](d_x, d_next)
        d_x = d_next
    return d_x.copy_to_host()[0]

rng = np.random.default_rng(0)
x = rng.random(2**24, dtype=np.float32)
d_x = cuda.to_device(x)
for launch in range(3):
    print(f"run {launch}: two-pass float32 total {total(d_x):.2f}")
print(f"NumPy float64 sum:          {x.sum(dtype=np.float64):.2f}")

$ python p4_two_pass.py
/usr/local/lib/python3.13/dist-packages/numba_cuda/numba/cuda/dispatcher.py:716: NumbaPerformanceWarning: Grid size 1 will likely result in GPU under-utilization due to low occupancy.
  warn(errors.NumbaPerformanceWarning(msg))
run 0: two-pass float32 total 8388153.50
run 1: two-pass float32 total 8388153.50
run 2: two-pass float32 total 8388153.50
NumPy float64 sum:          8388153.94