支持的原子操作

CUDA 内置目标弃用通知

Numba 内置的 CUDA 目标已弃用，后续开发已转移到 NVIDIA numba-cuda 软件包。请参阅 内置 CUDA 目标弃用和维护状态。

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 赋值给该元素。

仅支持 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))