原文: http://numba.pydata.org/numba-doc/latest/cuda/reduction.html
为 CUDA GPU 编写缩减算法可能很棘手。 Numba 提供了一个@reduce
装饰器,用于将简单的二进制操作转换为简化内核。
例:
import numpy
from numba import cuda
@cuda.reduce
def sum_reduce(a, b):
return a + b
A = (numpy.arange(1234, dtype=numpy.float64)) + 1
expect = A.sum() # numpy sum reduction
got = sum_reduce(A) # cuda sum reduction
assert expect == got
Lambda 函数也可以在这里使用:
sum_reduce = cuda.reduce(lambda a, b: a + b)
reduce
装饰器创建Reduce
类的实例。 (目前,reduce
是Reduce
的别名,但不保证这种行为。)
class numba.cuda.Reduce(functor)
__call__(arr, size=None, res=None, init=0, stream=0)
完全减少。
| 参数: |
- arr - 主机或设备阵列。如果给出了设备数组,则会在原地执行缩减,并覆盖数组中的值。如果给出了主机阵列,则会自动将其复制到设备。
- size - 可选整数,指定要减少的
arr
中的元素数。如果未指定此参数,则会减少整个数组。 - res - 将减少结果写入的可选设备数组。结果写入此数组的第一个元素。如果指定了此参数,则不会从设备到主机进行减少输出的通信。
- init - 还原的可选初始值,其类型必须与
arr.dtype
匹配。 - 流 - 执行还原的可选 CUDA 流。如果未指定流,则使用默认流 0。
|
| --- | --- |
| 返回: | 如果指定了res
,则返回None
。否则,返回减少的结果。 |
| --- | --- |
__init__(functor)
创建一个使用给定二进制函数减少值的缩减对象。二进制函数编译一次并缓存在此对象中。保持此对象存活将阻止重新编译。
参数: | binop - 要编译为 CUDA 设备函数的函数,该函数将用作 CUDA 设备上的二进制运算。在内部,它使用cuda.jit(device=True) 编译。 |
---|