Numba cuda：使用共享内存添加数字导致覆盖

Question

我一直在尝试使用共享内存添加数字，所以它如下所示：

线程0：将共享内存变量sharedMemT[0]加1

线程1：将共享内存变量sharedMemT[0]加1

同步线程并将 sharedMemT[0] 存储到输出[0]

但结果是... 1??

@cuda.jit()
def add(output):
    sharedMemT = cuda.shared.array(shape=(1), dtype=int32)
    sharedMemT[0] = 0
    cuda.syncthreads()

    sharedMemT[0] += 1
    cuda.syncthreads()
    output[0] = sharedMemT[0]

out = np.array([0])
add[1, 2](out)
print(out) # results in [1]

Answer 1

恭喜，您参加了记忆竞赛。线程0和1同时运行，所以结果是undefined，无论是对共享内存变量的操作，还是回写到全局内存中。

为了使其正常工作，您需要使用原子内存操作序列化对共享内存变量的访问，然后只有一个线程写回全局内存：

$ cat atomic.py

import numpy as np
from numba import cuda, int32

@cuda.jit()
def add(output):
    sharedMemT = cuda.shared.array(shape=(1), dtype=int32)
    pos = cuda.grid(1)
    if pos == 0:
        sharedMemT[0] = 0

    cuda.syncthreads()

    cuda.atomic.add(sharedMemT, 0, 1)
    cuda.syncthreads()

    if pos == 0:
        output[0] = sharedMemT[0]

out = np.array([0])
add[1, 2](out)
print(out)

$ python atomic.py
[2]

Numba cuda：使用共享内存添加数字导致覆盖

Numba cuda: Using shared memory to add numbers results in overwriting

python

cuda

shared-memory

numba