支持的原子操作

Numba 提供了对 CUDA 支持的一些原子操作的访问。目前实现的有以下几种：

class numba.cuda.atomic

原子操作的命名空间

class add(ary, idx, val)

执行原子操作 ary[idx] += val。仅支持 int32、float32 和 float64 操作数。

返回索引位置的旧值，就像它是以原子方式加载的一样。

class and_(ary, idx, val)

执行原子操作 ary[idx] &= val 。仅支持 int32、int64、uint32 和 uint64 操作数。

返回索引位置的旧值，就像它是以原子方式加载的一样。

class cas(ary, idx, old, val)

如果数组 ary 的元素 idx 的当前值与 old 匹配，则有条件地将 val 赋值给 ary 的元素 idx。

仅支持 int32、int64、uint32、uint64 操作数。

返回旧值，如同它是以原子方式加载的。

class compare_and_swap(ary, old, val)

如果当前值与 old 匹配，则有条件地将 val 赋值给一维数组 ary 的第一个元素。

仅支持 int32、int64、uint32、uint64 操作数。

返回旧值，如同它是以原子方式加载的。

class dec(ary, idx, val)

执行:

ary[idx] = (value if (array[idx] == 0) or
            (array[idx] > value) else array[idx] - 1)

仅支持在 uint32 和 uint64 操作数上使用。

返回索引位置的旧值，就像它是以原子方式加载的一样。

class exch(ary, idx, val)

执行原子操作 ary[idx] = val 。仅支持 int32、int64、uint32 和 uint64 操作数。

返回索引位置的旧值，就像它是以原子方式加载的一样。

class inc(ary, idx, val)

执行原子 ary[idx] += 1 操作，直到达到 val，然后重置为 0。仅支持 uint32 和 uint64 操作数。

返回索引位置的旧值，就像它是以原子方式加载的一样。

class max(ary, idx, val)

执行原子操作 ary[idx] = max(ary[idx], val)。

仅支持 int32、int64、uint32、uint64、float32、float64 操作数。

返回索引位置的旧值，就像它是以原子方式加载的一样。

class min(ary, idx, val)

执行原子操作 ary[idx] = min(ary[idx], val)。

仅支持 int32、int64、uint32、uint64、float32、float64 操作数。

返回索引位置的旧值，就像它是以原子方式加载的一样。

class nanmax(ary, idx, val)

执行原子操作 ary[idx] = max(ary[idx], val)。

注意：NaN 被视为缺失值，因此：nanmax(NaN, n) == n, nanmax(n, NaN) == n

仅支持 int32、int64、uint32、uint64、float32、float64 操作数。

返回索引位置的旧值，就像它是以原子方式加载的一样。

class nanmin(ary, idx, val)

执行原子操作 ary[idx] = min(ary[idx], val)。

注意：NaN 被视为缺失值，因此：nanmin(NaN, n) == n, nanmin(n, NaN) == n

仅支持 int32、int64、uint32、uint64、float32、float64 操作数。

返回索引位置的旧值，就像它是以原子方式加载的一样。

class or_(ary, idx, val)

执行原子操作 ary[idx] |= val 。仅支持 int32、int64、uint32 和 uint64 操作数。

返回索引位置的旧值，就像它是以原子方式加载的一样。

class sub(ary, idx, val)

执行原子操作 ary[idx] -= val 。仅支持 int32、float32 和 float64 操作数。

返回索引位置的旧值，就像它是以原子方式加载的一样。

class xor(ary, idx, val)

执行原子操作 ary[idx] ^= val。仅支持 int32、int64、uint32 和 uint64 操作数。

返回索引位置的旧值，就像它是以原子方式加载的一样。

示例

以下代码演示了如何使用 numba.cuda.atomic.max 在一个数组中找到最大值。注意，在这种情况下，这不是找到最大值的最有效方法，但它作为一个示例:

from numba import cuda
import numpy as np

@cuda.jit
def max_example(result, values):
    """Find the maximum value in values and store in result[0]"""
    tid = cuda.threadIdx.x
    bid = cuda.blockIdx.x
    bdim = cuda.blockDim.x
    i = (bid * bdim) + tid
    cuda.atomic.max(result, 0, values[i])


arr = np.random.rand(16384)
result = np.zeros(1, dtype=np.float64)

max_example[256,64](result, arr)
print(result[0]) # Found using cuda.atomic.max
print(max(arr))  # Print max(arr) for comparison (should be equal!)

多维数组通过使用一个整数元组作为索引来支持:

@cuda.jit
def max_example_3d(result, values):
    """
    Find the maximum value in values and store in result[0].
    Both result and values are 3d arrays.
    """
    i, j, k = cuda.grid(3)
    # Atomically store to result[0,1,2] from values[i, j, k]
    cuda.atomic.max(result, (0, 1, 2), values[i, j, k])

arr = np.random.rand(1000).reshape(10,10,10)
result = np.zeros((3, 3, 3), dtype=np.float64)
max_example_3d[(2, 2, 2), (5, 5, 5)](result, arr)
print(result[0, 1, 2], '==', np.max(arr))