[英]Numba - Shared memory in CUDA kernel not updating correctly
考慮以下 kernel,它計算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])
原子添加在相對較新的設備上進行了優化,例如使用硬件扭曲減少,但是當大量流式多處理器執行原子操作時,操作往往不會擴展。
提高此 kernel 性能的一種解決方案是在陣列足夠大的情況下執行許多值的塊縮減。 在實踐中,每個線程可以對多個項目求和,最后執行一個原子操作。 代碼應如下所示(未經測試):
# 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)和相對較小的 arrays 使用較小的值肯定會更快。
聲明:本站的技術帖子網頁,遵循CC BY-SA 4.0協議,如果您需要轉載,請注明本站網址或者原文地址。任何問題請咨詢:yoyou2525@163.com.