Cuda 并行化内核共享计数器变量

Cuda Parallelized Kernel Shared Counter Variable

有没有办法让一个整数计数器变量可以 incremented/decremented 跨并行化 cuda 内核中的所有线程?下面的代码输出“[1]”,因为一个线程对计数器数组的修改不会应用于其他线程。

import numpy as np
from numba import cuda


@cuda.jit('void(int32[:])')
def func(counter):
    counter[0] = counter[0] + 1


counter = cuda.to_device(np.zeros(1, dtype=np.int32))
threadsperblock = 64
blockspergrid = 18
func[blockspergrid, threadsperblock](counter)
print(counter.copy_to_host())

一种方法是使用 numba cuda atomics:

$ cat t18.py
import numpy as np
from numba import cuda


@cuda.jit('void(int32[:])')
def func(counter):
    cuda.atomic.add(counter, 0, 1)


counter = cuda.to_device(np.zeros(1, dtype=np.int32))
threadsperblock = 64
blockspergrid = 18
print blockspergrid * threadsperblock
func[blockspergrid, threadsperblock](counter)
print(counter.copy_to_host())
$ python t18.py
1152
[1152]
$

一个atomic operation对目标执行不可分割的读取-修改-写入操作,因此线程在更新目标变量时不会相互干扰。

当然也可以有其他的方法,根据自己的实际需要,比如classical parallel reduction. numba provides some reduction sugar也可以。