Numba

Question

考虑以下内核，它计算 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]

Answer 1

您需要显式执行原子添加操作:

@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])

原子添加在相对较新的设备上进行了优化，例如使用硬件扭曲减少，但当大量 streaming-multiprocessors 执行原子操作时，操作往往不会扩展。

提高此内核性能的一种解决方案是在假设数组足够大的情况下对许多值执行块缩减。实际上，每个线程可以对多个项目求和，最后执行一个原子操作。代码应如下所示（未经测试）：

# 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）和对于相对较小的数组使用较小的值肯定更快。

Numba - CUDA 内核中的共享内存未正确更新

Numba - Shared memory in CUDA kernel not updating correctly

python

cuda