原文: 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 操作数上受支持。
以下代码演示了如何使用 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))