考虑以下内核,它计算x
中小于或等于y
中相应元素的元素数。
@cuda.jit
def count_leq(x, y, out):
i = cuda.grid(1)
shared = cuda.shared.array(1, dtype=DTYPE)
if i < len(x):
shared[0] += x[i] <= y[i]
cuda.syncthreads()
out[0] = shared[0]
但是,每个线程的增量没有正确保存在共享数组中。
a = cuda.to_device(np.arange(5)) # [0 1 2 3 4]
b = cuda.to_device(np.arange(5)) # [0 1 2 3 4]
out = cuda.to_device(np.zeros(1)) # [0]
count_leq[1,len(a)](a, b, out)
print(out[0]) # 1.0, but should be 5.0
我在这里做错了什么?我很困惑,因为cuda.shared.array
是由给定块中的所有线程共享的,对吧?如何使用相同的1元素数组来累积增量?
我还尝试了以下操作,但以与上述版本相同的行为失败。
@cuda.jit
def count_leq(x, y, out):
i = cuda.grid(1)
if i < len(x):
out[0] += x[i] <= y[i]
您需要显式执行原子添加操作:
@cuda.jit
def count_leq(x, y, out):
i = cuda.grid(1)
if i < len(x):
cuda.atomic.add(out, 0, x[i] <= y[i])
原子添加在相对较新的设备上进行了优化,例如使用硬件扭曲减少,但当大量流式多处理器执行原子操作时,操作往往不会扩展。
提高该内核性能的一个解决方案是,假设数组足够大,则执行多个值的块缩减。在实践中,每个线程可以对多个项目求和,最后执行一个原子操作。代码应该是这样的(未经测试(:
# Must be launched with different parameters since
# each threads works on more array items.
# The number of block should be 16 times smaller.
@cuda.jit
def count_leq(x, y, out):
tid = cuda.threadIdx.x
bid = cuda.blockIdx.x
bdim = cuda.blockDim.x
i = (bid * bdim * 16) + tid
s = 0
# Fast general case (far from the end of the arrays)
if i+16*bdim < len(x):
# Thread-local reduction
# This loop should be unrolled
for j in range(16):
idx = i + j * bdim
s += x[idx] <= y[idx]
# Slower corner case (close to end of the arrays: checks are needed)
else:
for j in range(16):
idx = i + j * bdim
if idx < len(x):
s += x[idx] <= y[idx]
cuda.atomic.add(out, 0, s)
注意,16是任意值。对于较大的数组使用较大的值(如64(,而对于相对较小的数组使用较小的值,当然会更快。