3.6。支持的原子操作

原文： http://numba.pydata.org/numba-doc/latest/cuda/intrinsics.html

Numba 提供了对numba.cuda.atomic类中 CUDA 支持的一些原子操作的访问。

目前实施的内容如下：

class numba.cuda.atomic

用于原子操作的命名空间

class add(ary, idx, val)

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

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

class compare_and_swap(ary, old, val)

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

返回当前值，就像它以原子方式加载一样。

class max(ary, idx, val)

执行原子 ary [idx] = max（ary [idx]，val）。 NaN 被视为缺失值，因此 max（NaN，n）== max（n，NaN）== n。请注意，这与 Python 和 Numpy 行为不同，其中当 a 或 b 是 NaN 时，max（a，b）始终为 a。

仅在 int32，int64，uint32，uint64，float32，float64 操作数上受支持。

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

class min(ary, idx, val)

执行原子 ary [idx] = min（ary [idx]，val）。 NaN 被视为缺失值，因此 min（NaN，n）== min（n，NaN）== n。请注意，这与 Python 和 Numpy 行为不同，其中 min（a，b）始终是 a 或 b 是 NaN 时的行为。

仅在 int32，int64，uint32，uint64，float32，float64 操作数上受支持。

3.6.1。示例

以下代码演示了如何使用 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 comparision (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))