GPU 归约

为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 是 Reduce 的别名，但这种行为并不保证。

class numba.cuda.Reduce(functor)[源代码]

创建一个归约对象，该对象使用给定的二元函数来归约值。二元函数会被编译一次并缓存在此对象内部。保持此对象存活将防止重新编译。

__init__(functor)[源代码]

参数:: functor – 实现归约二元操作的函数。它将被编译为使用 cuda.jit(device=True) 的 CUDA 设备函数。

__call__(arr, size=None, res=None, init=0, stream=0)[源代码]

执行完全约简。

参数:

arr – 主机或设备数组。
size – 可选的整数，指定在 arr 中要减少的元素数量。如果未指定此参数，则整个数组将被减少。
res – 可选的设备数组，用于写入归约结果。结果被写入此数组的第一个元素。如果指定了此参数，则不会发生从设备到主机的归约输出通信。
init – 归约的可选初始值，其类型必须与 arr.dtype 匹配。
stream – 执行归约的可选 CUDA 流。如果没有指定流，则使用默认流 0。

返回:

如果指定了 res ，则返回 None 。否则，返回归约的结果。