使用 numba 原子操作函数时遇到问题 (cuda.atomic.compare_and_swap)

Question

我正在尝试使用 Numba 为我的代码编写 cuda 内核。不知何故，我想在我的部分代码中使用原子操作，我编写了一个测试内核来查看 cuda.atomic.compare_and_swap 是如何工作的。在文档中它是这样说的：

enter image description here

from numba import cuda
import numpy as np


@cuda.jit
def atomicCAS(N,out1):
    idx = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x
    if idx >= N:
        return
    A = out1[idx:]
    cuda.atomic.compare_and_swap(A,idx,0)

N    = 1024
out1 = np.arange(N)
out1 = np.zeros(N)
dout1 = cuda.to_device(out1)
tpb  = 32
bpg  = int(np.ceil(N/tpb))
atomicCAS[bpg,tpb](N,dout1)
hout1 = dout1.copy_to_host()

然后我得到这个错误：


TypingError: Invalid use of Function(<class 'numba.cuda.stubs.atomic.compare_and_swap'>) with argument(s) of type(s): (array(float64, 1d, A), int64, Literal[int](0))
 * parameterized
In definition 0:
    All templates rejected with literals.
In definition 1:
    All templates rejected without literals.
This error is usually caused by passing an argument of a type that is unsupported by the named function.
[1] During: resolving callee type: Function(<class 'numba.cuda.stubs.atomic.compare_and_swap'>)
[2] During: typing of call at /home/qinyu/test.py (20)

这是一个非常天真的代码，我想我输入了变量的写入类型，但我遇到了这个打字错误。它与 Numba 中的其他原子操作配合得很好。这是唯一对我不起作用的。有人可以帮我找出问题所在，还是有其他替代方法可以解决这个问题？谢谢！

Answer 1

报错信息中的关键是：

array(float64, 1d, A), int64, Literal[int](0))

仅限 CUDA atomicCAS supports integer types。您不能传递浮点类型。

使用 numba 原子操作函数时遇到问题 (cuda.atomic.compare_and_swap)

Have trouble using numba atomic operation functions (cuda.atomic.compare_and_swap)

python

cuda

atomic

compare-and-swap

numba