有没有办法让一个整数计数器变量可以在并行化的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也提供了一些reduction sugar。